Update app.py

app.py

@@ -1,126 +1,325 @@
 import streamlit as st
-from streamlit import session_state as session
+import spotipy
+from spotipy.oauth2 import SpotifyOAuth
+from qdrant_client import QdrantClient
+from qdrant_client.http import models
 from src.laion_clap.inference import AudioEncoder
-# from src.utils.spotify import SpotifyHandler, SpotifyAuthentication
-import pandas as pd
-from dotenv import load_dotenv
-from langchain.llms import CTransformers, Ollama
-from src.llm.chain import LLMChain
-from pymongo.mongo_client import MongoClient
 import os
+import re
+import unicodedata
+import requests
+import uuid
 
-st.set_page_config(page_title="Curate me a playlist", layout="wide")
-load_dotenv()
-
-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=os.getenv("LLM_VERSION"),
-        # 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
+# Spotify API credentials
+SPOTIPY_CLIENT_ID = 'd927d12613c4418d85313f69ce298987'
+SPOTIPY_CLIENT_SECRET = '95440b13cb3f466f922d0228290be4ff'
+SPOTIPY_REDIRECT_URI = 'http://localhost:8501/'
+SCOPE = 'user-library-read'
+CACHE_PATH = '.spotifycache'
+
+# Qdrant setup
+QDRANT_HOST = "localhost"
+QDRANT_PORT = 6333
+COLLECTION_NAME = "spotify_songs"
+
+st.set_page_config(page_title="Spotify Similarity Search", page_icon="🎵", layout="wide")
 
 @st.cache_resource
 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
+    return AudioEncoder()
 
 @st.cache_resource
-def output_songs(text):
-    output = llm_pipeline.process_user_description(text)
-    if output:
-        song_list = []
-        for _, song_desc in output:
-            print(song_desc)
-            ranking = recommender.list_top_k_songs(song_desc, k=15)
-            song_list += ranking
-        return  pd.DataFrame(song_list)\
-            .sort_values("score", ascending=False)\
-                .drop_duplicates(subset=["track_id"])\
-                    .reset_index(drop=True)
-    else:
+def get_qdrant_client():
+    client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
+    try:
+        client.get_collection(COLLECTION_NAME)
+    except Exception:
+        st.error("Qdrant collection not found. Please ensure the collection is properly initialized.")
+    return client
+
+def get_spotify_client():
+    auth_manager = SpotifyOAuth(
+        client_id=SPOTIPY_CLIENT_ID,
+        client_secret=SPOTIPY_CLIENT_SECRET,
+        redirect_uri=SPOTIPY_REDIRECT_URI,
+        scope=SCOPE,
+        cache_path=CACHE_PATH
+    )
+    
+    if 'code' in st.experimental_get_query_params():
+        token_info = auth_manager.get_access_token(st.experimental_get_query_params()['code'][0])
+        return spotipy.Spotify(auth=token_info['access_token'])
+    
+    if not auth_manager.get_cached_token():
+        auth_url = auth_manager.get_authorize_url()
+        st.markdown(f"[Click here to login with Spotify]({auth_url})")
         return None
+    
+    return spotipy.Spotify(auth_manager=auth_manager)
+
+def find_similar_songs_by_text(_query_text, _qdrant_client, _text_encoder, top_k=10):
+    query_vector = generate_text_embedding(_query_text, _text_encoder)
+    search_result = _qdrant_client.query_points(
+        collection_name=COLLECTION_NAME,
+        query=query_vector.tolist()[0],
+        limit=top_k
+    ).model_dump()["points"]
+    return [
+        {
+            "name": hit["payload"]["name"],
+            "artist": hit["payload"]["artists"][0]["name"],
+            "similarity": hit["score"],
+            "preview_url": hit["payload"]["preview_url"]
+        } for hit in search_result
+    ]
 
+def generate_text_embedding(text, text_encoder):
+    text_data = [text]
+    return text_encoder.get_text_embedding(text_data)
 
+def logout():
+    if os.path.exists(CACHE_PATH):
+        os.remove(CACHE_PATH)
+    for key in list(st.session_state.keys()):
+        del st.session_state[key]
+    st.experimental_rerun()
 
-recommender, llm_pipeline = load_resources()
-
-st.title("""Curate me a Playlist.""")
-st.info("""
-
-        Hey there, introducing the Music Playlist Curator AI! It's designed to craft playlists based on your descriptions. 
-        Here's the breakdown: we've got a Mistral 7B-Instruct 5-bit quantized version LLM running on the CPU to handle user inputs, 
-        and a Contrastive Learning model from the Amazing [LAION AI](https://github.com/LAION-AI/CLAP) team for Audio-Text joint embeddings, scoring song similarity.
-
-        The songs are pulled from my personal Spotify Liked Songs through API. Using an automated data extraction pipeline, 
-        I queried each song on my list on YouTube, downloaded it, 
-        extracted audio features, and stored them on MongoDB.
-
-        TODOs:
-        - [ ] Making playlists on users' own Spotify Tracks,
-        - [ ] Display leaderboard to show the best playlist curated,
-        - [ ] Generate the playlist on Spotify directly
-""")
-st.success("The pipeline running on CPU which might take a few minutes to process.")
-
-st.warning("""
-        A caveat: because the audio data is retrieved from YouTube, 
-        there's a chance some songs might not be top-notch quality or could be live versions, impacting the audio features' quality.
-           
-        Another caveat: I've given it a spin with some Turkish descriptions, had some wins and some misses. I might wanna upgrade to a GPU powered environment
-        to enchance LLM capacity in the future.
-        Give it a shot and see how it goes! 🎶
-           """)
-
-# st.success("""
-
-# """)
-
-session.text_input = st.text_input(label="Describe a playlist")
-session.slider_count = st.slider(label="How many tracks", min_value=5, max_value=35, step=5)
-buffer1, col1, buffer2 = st.columns([1.45, 1, 1])
-
-is_clicked = col1.button(label="Curate")
-if is_clicked:
-
-    dataframe = output_songs(session.text_input)
-    if isinstance(dataframe, pd.DataFrame):
-        dataframe = dataframe.iloc[:session.slider_count]
-        dataframe.drop_duplicates(subset=["track_id"], inplace=True)
-        dataframe.drop(columns=["track_id", "score"], inplace=True)
-        st.data_editor(
-            dataframe,
-            column_config={
-                "link": st.column_config.LinkColumn(
-                    "link",
-                        )
-                            },
-            hide_index=False,
-            use_container_width=True
+def truncate_qdrant_data(qdrant_client):
+    try:
+        qdrant_client.delete_collection(collection_name=COLLECTION_NAME)
+        qdrant_client.create_collection(
+            collection_name=COLLECTION_NAME,
+            vectors_config=models.VectorParams(size=512, distance=models.Distance.COSINE),
         )
+        st.success("Qdrant data has been truncated successfully.")
+    except Exception as e:
+        st.error(f"An error occurred while truncating Qdrant data: {str(e)}")
+
+@st.cache_data
+def fetch_all_liked_songs(_sp):
+    all_songs = []
+    offset = 0
+    while True:
+        results = _sp.current_user_saved_tracks(limit=50, offset=offset)
+        if not results['items']:
+            break
+        all_songs.extend([{
+            'id': item['track']['id'],
+            'name': item['track']['name'],
+            'artists': [{'name': artist['name'], 'id': artist['id']} for artist in item['track']['artists']],
+            'album': {
+                'name': item['track']['album']['name'],
+                'id': item['track']['album']['id'],
+                'release_date': item['track']['album']['release_date'],
+                'total_tracks': item['track']['album']['total_tracks']
+            },
+            'duration_ms': item['track']['duration_ms'],
+            'explicit': item['track']['explicit'],
+            'popularity': item['track']['popularity'],
+            'preview_url': item['track']['preview_url'],
+            'added_at': item['added_at'],
+            'is_local': item['track']['is_local']
+        } for item in results['items']])
+        offset += len(results['items'])
+    return all_songs
+
+def sanitize_filename(filename):
+    filename = re.sub(r'[<>:"/\\|?*]', '', filename)
+    filename = re.sub(r'[\s.]+', '_', filename)
+    filename = unicodedata.normalize('NFKD', filename).encode('ASCII', 'ignore').decode()
+    return filename[:100]
+
+def get_preview_filename(song):
+    safe_name = sanitize_filename(f"{song['name']}_{song['artists'][0]['name']}")
+    return f"{safe_name}.mp3"
+
+def download_preview(preview_url, song):
+    if not preview_url:
+        return False, None
+    
+    filename = get_preview_filename(song)
+    output_path = os.path.join("previews", filename)
+    
+    if os.path.exists(output_path):
+        return True, output_path
+    
+    response = requests.get(preview_url)
+    if response.status_code == 200:
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        with open(output_path, 'wb') as f:
+            f.write(response.content)
+        return True, output_path
+    return False, None
+
+def process_song(song, audio_encoder, qdrant_client):
+    filename = get_preview_filename(song)
+    output_path = os.path.join("previews", filename)
+    
+    if os.path.exists(output_path):
+        return output_path, None
+
+    preview_url = song['preview_url']
+    
+    if not preview_url:
+        return None, f"No preview available for: {song['name']} by {song['artists'][0]['name']}"
+
+    success, file_path = download_preview(preview_url, song)
+    if success:
+        # Check if the song is already in Qdrant
+        existing_points = qdrant_client.scroll(
+            collection_name=COLLECTION_NAME,
+            scroll_filter=models.Filter(
+                must=[
+                    models.FieldCondition(
+                        key="spotify_id",
+                        match=models.MatchValue(value=song['id'])
+                    )
+                ]
+            ),
+            limit=1
+        )[0]
+
+        if not existing_points:
+            embedding = generate_audio_embedding(file_path, audio_encoder)
+            point_id = str(uuid.uuid4())
+            
+            qdrant_client.upsert(
+                collection_name=COLLECTION_NAME,
+                points=[
+                    models.PointStruct(
+                        id=point_id,
+                        vector=embedding,
+                        payload={
+                            "name": song['name'],
+                            "artists": song['artists'],
+                            "spotify_id": song['id'],
+                            "album": song['album'],
+                            "duration_ms": song['duration_ms'],
+                            "popularity": song['popularity'],
+                            "preview_url": song['preview_url'],
+                            "local_preview_path": file_path
+                        }
+                    )
+                ]
+            )
+        return file_path, None
     else:
-        st.warning("User prompt could not be processed.")
+        return None, f"Failed to download preview for: {song['name']} by {song['artists'][0]['name']}"
+
+
+def generate_audio_embedding(audio_path, audio_encoder):
+    # This is a placeholder. You'll need to implement the actual audio embedding generation
+    # based on how your audio_encoder works with local audio files
+    return audio_encoder.extract_audio_representaion(audio_path).tolist()[0]
+
+def retrieve_all_previews(sp, qdrant_client, audio_encoder):
+    all_songs = fetch_all_liked_songs(sp)
+    total_songs = len(all_songs)
+    
+    progress_bar = st.progress(0)
+    status_text = st.empty()
+    warnings = []
+
+    for i, song in enumerate(all_songs):
+        _, warning = process_song(song, audio_encoder, qdrant_client)
+        if warning:
+            warnings.append(warning)
+        
+        # Update progress
+        progress = (i + 1) / total_songs
+        progress_bar.progress(progress)
+        status_text.text(f"Processing: {i+1}/{total_songs} songs")
+    
+    st.success(f"Processed {total_songs} songs.")
+    return warnings
+
+def display_warnings(warnings):
+    if warnings:
+        with st.expander("Processing Warnings", expanded=False):
+            st.markdown("""
+            <style>
+            .warning-box {
+                background-color: #fff3cd;
+                border-left: 6px solid #ffeeba;
+                margin-bottom: 10px;
+                padding: 10px;
+                color: #856404;
+            }
+            </style>
+            """, unsafe_allow_html=True)
+            
+            for warning in warnings:
+                st.markdown(f'<div class="warning-box">{warning}</div>', unsafe_allow_html=True)
+
+def main():
+    
+    st.title("Spotify Similarity Search")
+
+    audio_encoder = load_resources()
+    qdrant_client = get_qdrant_client()
+    
+    # Sidebar for authentication and data management
+    with st.sidebar:
+        st.header("Authentication & Data Management")
+        if 'spotify_auth' not in st.session_state:
+            sp = get_spotify_client()
+            if sp:
+                st.session_state['spotify_auth'] = sp
+        
+        if 'spotify_auth' in st.session_state:
+            st.success("Connected to Spotify and Qdrant")
+            if st.button("Logout from Spotify"):
+                logout()
+            if st.button("Truncate Qdrant Data"):
+                truncate_qdrant_data(qdrant_client)
+            if st.button("Retrieve All Previews"):
+                with st.spinner("Retrieving previews..."):
+                    warnings = retrieve_all_previews(st.session_state['spotify_auth'], qdrant_client, audio_encoder)
+                display_warnings(warnings)
+        elif 'code' in st.experimental_get_query_params():
+            st.warning("Authentication in progress. Please refresh this page.")
+        else:
+            st.info("Please log in to access your Spotify data.")
+     # Main content area
+    if 'spotify_auth' in st.session_state:
+        # Quick Start Guide
+        st.info("""
+        ### 🚀 Quick Start Guide
+
+        1. 🔄 Click 'Retrieve All Previews' in the sidebar, to start getting 30 seconds raw audio previews.
+        2. 🔍 Enter descriptive keywords (e.g., "upbeat electronic with female vocals")
+        3. 🎵 Explore similar songs and enjoy!
+
+        Note: Some songs may not have previews available mainly due to Spotify restrictions.      
+        
+        ✅ Do: Use specific terms (genre, mood, instruments)
+                
+        ❌ Don't: Use artist names or song titles
+        
+        💡 Tip: Refine your search if results aren't perfect!
+        """)
 
+        st.header("Find Similar Songs")
+        query_text = st.text_input("Enter a description or keywords for the music you're looking for:")
+        
+        if st.button("Search Similar Songs") or query_text:
+            if query_text:
+                with st.spinner("Searching for similar songs..."):
+                    search_results = find_similar_songs_by_text(query_text, qdrant_client, audio_encoder)
+                
+                if search_results:
+                    st.subheader("Similar songs based on your description:")
+                    for song in search_results:
+                        st.write(f"{song['name']} by {song['artist']} (Similarity: {song['similarity']:.2f})")
+                        if song['preview_url']:
+                            st.audio(song['preview_url'], format='audio/mp3')
+                        else:
+                            st.write("No preview available")
+                        st.write("---")  # Add a separator between songs
+                else:
+                    st.info("No similar songs found. Try a different description.")
+            else:
+                st.warning("Please enter a description or keywords for your search.")
 
-    # with st.form(key="spotiform"):
-    #     st.form_submit_button(on_click=authenticate_spotify, args=(session.access_url, ))
-    #     st.markdown(session.access_url)
+if __name__ == "__main__":
+    main()