llm integration

app.py

@@ -1,79 +1,67 @@
 import streamlit as st
 from streamlit import session_state as session
-from src.config.configs import ProjectPaths
-import numpy as np
 from src.laion_clap.inference import AudioEncoder
-import pickle
-import torch
+from src.utils.spotify import SpotifyHandler, SpotifyAuthentication
 import pandas as pd
-import json
-import os
-import smtplib, ssl
 from dotenv import load_dotenv
-st.set_page_config(page_title="Curate me a playlist", layout="wide")
+from langchain.llms import CTransformers, Ollama
+from src.llm.chain import LLMChain
+from pymongo.mongo_client import MongoClient
+import os
 
+st.set_page_config(page_title="Curate me a playlist", layout="wide")
 load_dotenv()
 
-@st.cache_data
-def load_data():
-    vectors = np.load(ProjectPaths.DATA_DIR.joinpath("vectors", "audio_representations.npy"))
-    with open(ProjectPaths.DATA_DIR.joinpath("vectors", "song_names.pkl"), "rb") as reader:
-        song_names = pickle.load(reader)
-
-    with open(ProjectPaths.DATA_DIR.joinpath("json", "youtube_data.json"), "r") as reader:
-        youtube_data = json.load(reader)
-
-    df_youtube = pd.DataFrame(youtube_data)
-    df_youtube["id"] = df_youtube["artist_name"] + " - " + df_youtube["track_name"] + ".wav"
-    df_youtube.set_index("id", inplace=True)
-    return vectors, song_names, df_youtube
-
+def load_llm_pipeline():
+    ctransformers_config = {
+        "max_new_tokens": 3000,
+        "temperature": 0,
+        "top_k": 1,
+        "top_p": 1,
+        "context_length": 2800
+        }
+
+    llm = CTransformers(
+        model="TheBloke/Mistral-7B-Instruct-v0.1-GGUF", 
+        model_file="mistral-7b-instruct-v0.1.Q5_K_M.gguf",
+        config=ctransformers_config
+        )
+    # llm = Ollama(temperature=0, model="mistral:7b-instruct-q8_0", top_k=1, top_p=1, num_ctx=2800)
+    chain = LLMChain(llm)
+    return chain
 
 @st.cache_resource
-def load_model():
-    recommender = AudioEncoder()
-    return recommender
+def load_resources():
+    password = os.getenv("MONGODB_PASSWORD")
+    url = os.getenv("MONGODB_URL")
+    uri = f"mongodb+srv://berkaygkv:{password}@{url}/?retryWrites=true&w=majority"
+    client = MongoClient(uri)
+    db = client.spoti
+    mongo_db_collection = db.saved_tracks
+    recommender = AudioEncoder(mongo_db_collection)
+    recommender.load_existing_audio_vectors()
+    llm_pipeline = load_llm_pipeline()
+    return recommender, llm_pipeline
 
 
-def send_curator(text):
-    port = int(os.getenv("PORT"))
-    print(port)
-    smtp_server = "smtp.gmail.com"
-    sender_email = os.getenv("EMAIL_ADDRESS")
-    receiver_email = os.getenv("EMAIL_RECEIVER")
-    password = os.getenv("EMAIL_PASSWORD")
-    from email.mime.multipart import MIMEMultipart
-    from email.mime.text import MIMEText
-
-    msg = MIMEMultipart("alternative")
-    msg["Subject"] = "Curate me a playlist submission"
-    part1 = MIMEText(body, "plain")
-    msg.attach(part1)
-    context = ssl.create_default_context()
-    with smtplib.SMTP_SSL(smtp_server, port, context=context) as server:
-        server.login(sender_email, password)
-        server.sendmail(sender_email, receiver_email, msg)
-
-    print("Email sent.")
-
-
-recommender = load_model()
-audio_vectors, song_names, df_youtube = load_data()
+recommender, llm_pipeline = load_resources()
 
 st.title("""Curate me a Playlist.""")
 session.text_input = st.text_input(label="Describe a playlist")
-session.slider_count = st.slider(label="Track counts", min_value=5, max_value=30, step=5)
-buffer1, col1, buffer2 = st.columns([1.45, 1, 1])
+session.slider_count = st.slider(label="Track counts", min_value=5, max_value=35, step=5)
+buffer1, col1, col2, buffer2 = st.columns([1.45, 1, 1, 1])
 
 is_clicked = col1.button(label="Curate")
 if is_clicked:
-    text_embed = recommender.get_text_embedding(session.text_input)
-    with torch.no_grad():
-        ranking = torch.tensor(audio_vectors) @ torch.tensor(text_embed).t()
-        ranking = ranking[:, 0].reshape(-1, 1)
-    dataframe = pd.DataFrame(ranking, columns=[session.text_input], index=song_names).rename(columns={session.text_input: "score"})
-    dataframe = dataframe.merge(df_youtube[["link"]], left_index=True, right_index=True, how="left").nlargest(int(session.slider_count), "score")
-    # st.dataframe(dataframe, use_container_width=True)
+    output = llm_pipeline.process_user_description(session.text_input)
+    song_list = []
+    for _, song_desc in output:
+        print(song_desc)
+        ranking = recommender.list_top_k_songs(song_desc, k=15)
+        song_list += ranking
+    
+    dataframe = pd.DataFrame(song_list).sort_values("score", ascending=False).drop_duplicates(subset=["track_id"]).drop(columns=["track_id"]).reset_index(drop=True)
+    dataframe = dataframe.iloc[:session.slider_count]
     st.data_editor(
         dataframe,
         column_config={
@@ -88,22 +76,7 @@ if is_clicked:
         use_container_width=True
     )
 
-    form = st.form("form")
-    form.write("You can submit the playlist you've curated")
-    sender = form.text_input("Name of the curator")
-    query = session.text_input
-    playlist = [f"{k}\n" for k in dataframe.index]
-    playlist_string = "\n".join(dataframe.index.tolist())
-    body = f"""\
-    Subject: Curate me a playlist submission
-
-    Curator --> {sender}
-    Query --> {session.text_input}
-
-    Playlist
-    {playlist_string}
-    """
-
 
-    print(body)
-    is_submit = form.form_submit_button("Submit", on_click=send_curator, args=([body]))
+    # with st.form(key="spotiform"):
+    #     st.form_submit_button(on_click=authenticate_spotify, args=(session.access_url, ))
+    #     st.markdown(session.access_url)

src/laion_clap/inference.py

@@ -1,41 +1,102 @@
-import numpy as np
+from tqdm import tqdm
 import librosa
 import torch
 from src import laion_clap
-from glob import glob
-import pandas as pd
+import json
+import jmespath
 from ..config.configs import ProjectPaths
-import pickle
 
 
 class AudioEncoder(laion_clap.CLAP_Module):
-    def __init__(self) -> None:
-        super().__init__(enable_fusion=False, amodel='HTSAT-base')
+    def __init__(self, collection=None) -> None:
+        super().__init__(enable_fusion=False, amodel="HTSAT-base")
+        self.music_data = None
         self.load_ckpt(ckpt=ProjectPaths.MODEL_PATH)
+        self.collection = collection
+
+    # def _get_track_data(self):
+    #     with open(ProjectPaths.DATA_DIR.joinpath("json", "final_track_data.json"), "r") as reader:
+    #         track_data = json.load(reader)
+    #     return track_data
+
+    def _get_track_data(self):
+        data = self.collection.find({})
+        return data
+    
+
+    def update_collection_item(self, track_id, vector):
+        self.collection.update_one({"track_id": track_id}, {"$set": {"embedding": vector}})
+
 
     def extract_audio_representaion(self, file_name):
         audio_data, _ = librosa.load(file_name, sr=48000)
         audio_data = audio_data.reshape(1, -1)
+        audio_data = torch.from_numpy(audio_data)
         with torch.no_grad():
-            audio_embed = self.get_audio_embedding_from_data(x=audio_data, use_tensor=False)
+            audio_embed = self.get_audio_embedding_from_data(
+                x=audio_data, use_tensor=True
+            )
         return audio_embed
 
     def extract_bulk_audio_representaions(self, save=False):
-        music_files = glob(str(ProjectPaths.DATA_DIR.joinpath("audio", "*.wav")))
-        song_names = [k.split("/")[-1] for k in music_files]
-        music_data = np.zeros((len(music_files), 512), dtype=np.float32)
-        for m in range(music_data.shape[0]):
-            music_data[m] = self.extract_audio_representaion(music_files[m])
+        track_data = self._get_track_data()
+        processed_data = []
+        idx = 0
+        for track in tqdm(track_data):
+            if track["youtube_data"]["file_path"] and track["youtube_data"]["link"] not in processed_data:
+                tensor = self.extract_audio_representaion(track["youtube_data"]["file_path"])
+                self.update_collection_item(track["track_id"], tensor.tolist())
+                idx += 1
+
 
-        if not save:
-            return music_data, song_names
+    # def load_existing_audio_vectors(self):
+    #     self.music_data = torch.load(
+    #         ProjectPaths.DATA_DIR.joinpath("vectors", "audio_representations.pt")
+    #     )
+    #     with open(
+    #         ProjectPaths.DATA_DIR.joinpath("vectors", "final_track_data_w_links.json"),
+    #         "r",
+    #     ) as rd:
+    #         self.track_data = json.load(rd)
+
+    def load_existing_audio_vectors(self):
+        # embedding_result = list(self.collection.find({}, {"embedding": 1}))
+        # tracking_result = list(self.collection.find({}, {"embedding": 0}))
+        arrays = []
+        track_data = []
+        for idx, track in enumerate(self.collection.find({})):
+            if not track.get("embedding"):
+                continue
+            data = track.copy()
+            data.pop("embedding")
+            data.update({"vector_idx": idx})
+            arrays.append(track["embedding"][0])
+            track_data.append(data)
+
+        self.music_data = torch.tensor(arrays)
+        self.track_data = track_data.copy()
 
-        else:
-            np.save(ProjectPaths.DATA_DIR.joinpath("vectors", "audio_representations.npy"))
-            with open(ProjectPaths.DATA_DIR.joinpath("vectors", "song_names.pkl", "rb")) as writer:
-                pickle.dump(song_names, writer)
 
     def extract_text_representation(self, text):
         text_data = [text]
         text_embed = self.get_text_embedding(text_data)
         return text_embed
+
+    def list_top_k_songs(self, text, k=10):
+        assert self.music_data is not None
+        with torch.no_grad():
+            text_embed = self.get_text_embedding(text, use_tensor=True)
+
+        dot_product = self.music_data @ text_embed.T
+        top_10 = torch.topk(dot_product.flatten(), k)
+        indices = top_10.indices.tolist()
+        final_result = []
+        for k, i in enumerate(indices):
+            piece = {
+                "title": self.track_data[i]["youtube_data"]["title"],
+                "score": round(top_10.values[k].item(), 2),
+                "link": self.track_data[i]["youtube_data"]["link"],
+                "track_id": self.track_data[i]["track_id"],
+                }
+            final_result.append(piece)
+        return final_result

src/llm/chain.py

@@ -0,0 +1,51 @@
+from langchain.prompts import ChatPromptTemplate, PromptTemplate
+from langchain.schema.runnable import RunnableLambda
+from operator import itemgetter
+from langchain.output_parsers import PydanticOutputParser
+from .output_parser import SongDescriptions
+from langchain.llms.base import LLM
+import json
+
+
+class LLMChain:
+    def __init__(self, llm_model: LLM) -> None:
+        self.llm_model = llm_model
+        self.parser = PydanticOutputParser(pydantic_object=SongDescriptions)
+        self.full_chain = self._create_llm_chain()
+
+
+    def _get_output_format(self, _):
+        return self.parser.get_format_instructions()
+        
+    def _create_llm_chain(self):
+        prompt_response = ChatPromptTemplate.from_messages([
+            ("system", "You are an AI assistant, helping the user to turn a music playlist text description into four separate song descriptions that are probably contained in the playlist. Try to be specific with descriptions. Make sure all 4 song descriptions are similar.\n"),
+            ("system", "{format_instructions}\n"),
+            ("human", "Playlist description: {description}.\n"),
+            # ("human", "Song descriptions:"),
+        ])
+        # prompt = PromptTemplate(
+        #     template="You are an AI assistant, helping the user to turn a music playlist text description into three separate generic song descriptions that are probably contained in the playlist.\n{format_instructions}\n{description}\n",
+        #     input_variables=["description"],
+        #     partial_variables={"format_instructions": self.parser.get_format_instructions()},
+        # )
+
+
+        full_chain = (
+            {
+                "format_instructions": RunnableLambda(self._get_output_format),
+                "description": itemgetter("description"),
+            } 
+            | prompt_response
+            | self.llm_model
+        )
+        return full_chain
+    
+    def process_user_description(self, user_input):
+        output = self.full_chain.invoke(
+            {
+                "description": user_input
+            }
+        ).replace("\\", '')
+        return self.parser.parse(output)
+    

src/llm/output_parser.py

@@ -0,0 +1,8 @@
+from pydantic import BaseModel, Field
+
+
+class SongDescriptions(BaseModel):
+    song_description_1: str = Field(description="description of the first song")
+    song_description_2: str = Field(description="description of the second song")
+    song_description_3: str = Field(description="description of the third song")
+    song_description_4: str = Field(description="description of the fourth song")