feat(rec): Finalize recommendations (almost done)

.gitignore

@@ -1,4 +1,5 @@
 *__pycache__*
 .venv
 .env
-data
+data
+*sandbox*

Dockerfile

@@ -31,7 +31,7 @@ ENV HOME=/home/user \
 
 # Setup application directory
 WORKDIR $HOME/app
-ADD --chown=user ./prompts $HOME/app/prompts
+ADD --chown=user ./resources $HOME/app/resources
 ADD --chown=user ./ingest.py $HOME/app/ingest.py
 ADD --chown=user ./app.py $HOME/app/app.py
 

app.py

@@ -1,15 +1,25 @@
 import json
 import os
 import random
+import urllib.parse
 from pathlib import Path
+from typing import Optional
 
 import gradio as gr
+import numpy as np
+import pandas as pd
+from dotenv import load_dotenv
+from fastembed import SparseEmbedding, SparseTextEmbedding
 from google import genai
 from google.genai import types
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
+from qdrant_client import QdrantClient
+from qdrant_client import models as qmodels
+from sentence_transformers import CrossEncoder, SentenceTransformer
 from vllm import LLM, SamplingParams
 from vllm.sampling_params import GuidedDecodingParams
 
+load_dotenv()
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 
@@ -20,11 +30,44 @@ VLLM_DTYPE = os.getenv("VLLM_DTYPE")
 
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 
+DATA_PATH = Path(os.getenv("DATA_PATH"))
+DB_PATH = DATA_PATH / "db"
+
+client = QdrantClient(path=str(DB_PATH))
+collection_name = "knowledge_cards"
+num_chunks_base = 500
+alpha = 0.5
+top_k = 5  # we only want top 5 genres
+
+youtube_url_template = "{genre} music playlist"
+
+
 # -------------------------------- HELPERS -------------------------------------
-def load_prompt(path: Path) -> str:
+def load_text_resource(path: Path) -> str:
     with path.open("r") as file:
-        prompt = file.read()
-    return prompt
+        resource = file.read()
+    return resource
+
+
+def youtube_search_link_for_genre(genre: str) -> str:
+    base_url = "https://www.youtube.com/results"
+    params = {
+        "search_query": youtube_url_template.format(
+            genre=genre.replace("_", " ").lower()
+        )
+    }
+    return f"{base_url}?{urllib.parse.urlencode(params)}"
+
+
+def generate_recommendation_string(ranking: dict[str, float]) -> str:
+    recommendation_string = "## Recommendations for You\n\n"
+    for idx, (genre, score) in enumerate(ranking.items(), start=1):
+        youtube_link = youtube_search_link_for_genre(genre=genre)
+        recommendation_string += (
+            f"{idx}. **{genre.replace('_', ' ').capitalize()}** ({score:.2f}); "
+            f"[YouTube link]({youtube_link})\n"
+        )
+    return recommendation_string
 
 
 # --------------------------------  Data Models  -------------------------------
@@ -41,8 +84,51 @@ class QueryRewrite(BaseModel):
     structured: StructuredQueryRewriteResponse | None = None
 
 
+class APIGenreRecommendation(BaseModel):
+    name: str = Field(description="Name of the music genre.")
+    score: float = Field(
+        description="Score you assign to the genre (from 0 to 1).", ge=0, le=1
+    )
+
+
 class APIGenreRecommendationResponse(BaseModel):
-    genres: list[str]
+    genres: list[APIGenreRecommendation]
+
+
+class RetrievalResult(BaseModel):
+    chunk: str
+    genre: str
+    score: float
+
+
+class RerankingResult(BaseModel):
+    query: str
+    genre: str
+    chunk: str
+    score: float
+
+
+class Recommendation(BaseModel):
+    name: str
+    rank: int
+    score: Optional[float] = None
+
+
+class PipelineResult(BaseModel):
+    query: str
+    rewrite: Optional[QueryRewrite] = None
+    retrieval_result: Optional[list[RetrievalResult]] = None
+    reranking_result: Optional[list[RerankingResult]] = None
+    recommendations: Optional[dict[str, Recommendation]] = None
+
+    def to_ranking(self) -> dict[str, float]:
+        if not self.recommendations:
+            return {}
+
+        return {
+            genre: recommendation.score
+            for genre, recommendation in self.recommendations.items()
+        }
 
 
 # --------------------------------  VLLM  --------------------------------------
@@ -68,7 +154,7 @@ vllm_system_prompt = (
     "You are a search query optimization assistant built into"
     " music genre search engine, helping users discover novel music genres."
 )
-vllm_prompt = load_prompt(Path("./prompts/local.txt"))
+vllm_prompt = load_text_resource(Path("./resources/prompt_vllm.md"))
 
 # --------------------------------  GEMINI  ------------------------------------
 gemini_config = types.GenerateContentConfig(
@@ -76,20 +162,35 @@ gemini_config = types.GenerateContentConfig(
     response_schema=APIGenreRecommendationResponse,
     temperature=0.7,
     max_output_tokens=1024,
-    system_instruction=("You are a helpful music genre recommendation assistant."),
+    system_instruction=(
+        "You are a helpful music genre recommendation assistant built into"
+        " music genre search engine, helping users discover novel music genres."
+    )
 )
 gemini_llm = genai.Client(
     api_key=GEMINI_API_KEY,
     http_options={"api_version": "v1alpha"},
 )
-gemini_prompt = load_prompt(Path("./prompts/api.txt"))
+gemini_prompt = load_text_resource(Path("./resources/prompt_api.md"))
 
-
-# ---------------------------- RETRIEVAL ---------------------------------------
+# ---------------------------- EMBEDDING MODELS --------------------------------
+dense_encoder = SentenceTransformer(
+    model_name_or_path="mixedbread-ai/mxbai-embed-large-v1",
+    device="cuda",
+    model_kwargs={"torch_dtype": VLLM_DTYPE},
+)
+sparse_encoder = SparseTextEmbedding(model_name="Qdrant/bm25", cuda=True)
+reranker = CrossEncoder(
+    model_name_or_path="BAAI/bge-reranker-v2-m3",
+    max_length=1024,
+    device="cuda",
+    model_kwargs={"torch_dtype": VLLM_DTYPE},
+)
+reranker_batch_size = 128
 
 
-# ----------------------- GENERATE RECOMMENDATIONS -----------------------------
-def recommend_sadaimrec(query: str):
+# ---------------------------- RETRIEVAL ---------------------------------------
+def run_query_rewrite(query: str) -> QueryRewrite:
     prompt = vllm_prompt.format(query=query)
     messages = [
         {"role": "system", "content": vllm_system_prompt},
@@ -104,10 +205,181 @@ def recommend_sadaimrec(query: str):
         rewrites=[x for x in list(rewrite_json.values()) if x is not None],
         structured=rewrite_json,
     )
-    return f"SADAIMREC: response to '{rewrite.model_dump_json(indent=4)}'"
+    return rewrite
+
+
+def prepare_queries_for_retrieval(
+    query: str, rewrite: QueryRewrite
+) -> list[dict[str, str | None]]:
+    queries_to_retrieve = [{"text": query, "topic": None}]
+    for cat, rewrite in rewrite.structured.model_dump().items():
+        if rewrite is None:
+            continue
+        topic = cat
+        if cat not in ["subjective", "purpose", "technical"]:
+            topic = None
+        queries_to_retrieve.append({"text": rewrite, "topic": topic})
+    return queries_to_retrieve
+
+
+def run_retrieval(
+    queries: list[dict[str, str]],
+) -> RetrievalResult:
+    queries_to_embed = [query["text"] for query in queries]
+    dense_queries = list(
+        dense_encoder.encode(
+            queries_to_embed, convert_to_numpy=True, normalize_embeddings=True
+        )
+    )
+    sparse_queries = list(sparse_encoder.query_embed(queries_to_embed))
+    prefetches: list[qmodels.Prefetch] = []
+
+    for query, dense_query, sparse_query in zip(queries, dense_queries, sparse_queries):
+        assert dense_query is not None and sparse_query is not None
+        assert isinstance(dense_query, np.ndarray) and isinstance(
+            sparse_query, SparseEmbedding
+        )
+        topic = query.get("topic", None)
+        prefetch = [
+            qmodels.Prefetch(
+                query=dense_query,
+                using="dense",
+                filter=qmodels.Filter(
+                    must=[
+                        qmodels.FieldCondition(
+                            key="topic", match=qmodels.MatchValue(value=topic)
+                        )
+                    ]
+                )
+                if topic is not None
+                else None,
+                limit=num_chunks_base,
+            ),
+            qmodels.Prefetch(
+                query=qmodels.SparseVector(**sparse_query.as_object()),
+                using="sparse",
+                filter=qmodels.Filter(
+                    must=[
+                        qmodels.FieldCondition(
+                            key="topic", match=qmodels.MatchValue(value=topic)
+                        )
+                    ]
+                )
+                if topic is not None
+                else None,
+                limit=num_chunks_base,
+            ),
+        ]
+        prefetches.extend(prefetch)
+
+    retrieval_results = client.query_points(
+        collection_name=collection_name,
+        prefetch=prefetches,
+        query=qmodels.FusionQuery(fusion=qmodels.Fusion.RRF),
+        limit=num_chunks_base,
+    )
+
+    final_hits: list[RetrievalResult] = [
+        RetrievalResult(
+            chunk=hit.payload["text"], genre=hit.payload["genre"], score=hit.score
+        )
+        for hit in retrieval_results.points
+    ]
+    return final_hits
+
+
+def run_reranking(
+    query: str, retrieval_result: list[RetrievalResult]
+) -> list[RerankingResult]:
+    hit_texts: list[str] = [result.chunk for result in retrieval_result]
+    hit_genres: list[str] = [result.genre for result in retrieval_result]
+    hit_rerank = reranker.rank(
+        query=query,
+        documents=hit_texts,
+        batch_size=reranker_batch_size,
+    )
+    ranking = [
+        RerankingResult(
+            query=query,
+            genre=hit_genres[hit["corpus_id"]],
+            chunk=hit_texts[hit["corpus_id"]],
+            score=hit["score"],
+        )
+        for hit in hit_rerank
+    ]
+    ranking.sort(key=lambda x: x.score, reverse=True)
+    return ranking
+
+
+def get_top_genres(
+    df: pd.DataFrame,
+    column: str,
+    alpha: float = 1.0,
+    # beta: float = 1.0,
+    top_k: int | None = None,
+) -> pd.Series:
+    assert 0 <= alpha <= 1.0
+
+    # Min-max normalization of re-ranker scores before aggregation
+    task_scores = df[column]
+    min_score = task_scores.min()
+    max_score = task_scores.max()
+    if max_score > min_score:  # Avoid division by zero
+        df.loc[:, column] = (task_scores - min_score) / (max_score - min_score)
+
+    tg_df = df.groupby("genre").agg(size=("chunk", "size"), score=(column, "sum"))
+    tg_df["weighted_score"] = alpha * (tg_df["size"] / tg_df["size"].max()) + (
+        1 - alpha
+    ) * (tg_df["score"] / tg_df["score"].max())
+    tg = tg_df.sort_values("weighted_score", ascending=False)["weighted_score"]
+
+    if top_k:
+        tg = tg.head(top_k)
+
+    return tg
+
+
+def get_recommendations(
+    reranking_result: list[RerankingResult],
+) -> dict[str, Recommendation]:
+    ranking_df = pd.DataFrame([x.model_dump(mode="python") for x in reranking_result])
+    top_genres_series = get_top_genres(
+        df=ranking_df, column="score", alpha=alpha, top_k=top_k
+    )
+    recommendations = {
+        genre: Recommendation(name=genre, rank=rank, score=score)
+        for rank, (genre, score) in enumerate(
+            top_genres_series.to_dict().items(), start=1
+        )
+    }
+    return recommendations
+
+
+# ----------------------- GENERATE RECOMMENDATIONS -----------------------------
+def recommend_sadaimrec(query: str):
+    result = PipelineResult(query=query)
+    print("Running query processing...", flush=True)
+    result.rewrite = run_query_rewrite(query=query)
+    queries_to_retrieve = prepare_queries_for_retrieval(
+        query=query, rewrite=result.rewrite
+    )
+
+    print("Running retrieval...", flush=True)
+    result.retrieval_result = run_retrieval(queries_to_retrieve)
+
+    print("Running re-ranking...", flush=True)
+    result.reranking_result = run_reranking(
+        query=query, retrieval_result=result.retrieval_result
+    )
+
+    print("Aggregating recommendations...", flush=True)
+    result.recommendations = get_recommendations(result.reranking_result)
+    recommendation_string = generate_recommendation_string(result.to_ranking())
+    return f"{recommendation_string}"
 
 
 def recommend_gemini(query: str):
+    print("Generating recommendations using Gemini...", flush=True)
     prompt = gemini_prompt.format(query=query)
     response = gemini_llm.models.generate_content(
         model="gemini-2.0-flash",
@@ -115,17 +387,19 @@ def recommend_gemini(query: str):
         config=gemini_config,
     )
     parsed_content: APIGenreRecommendationResponse = response.parsed
-    return f"CHATGPT: response to '{parsed_content.model_dump_json(indent=4)}'"
+    parsed_content.genres.sort(key=lambda x: x.score, reverse=True)
+    ranking = {x.name.lower(): x.score for x in parsed_content.genres}
+    recommendation_string = generate_recommendation_string(ranking)
+    return f"{recommendation_string}"
 
 
-# Mapping names to functions
+# -------------------------------------- INTERFACE -----------------------------
 pipelines = {
     "sadaimrec": recommend_sadaimrec,
     "chatgpt": recommend_gemini,
 }
 
 
-# -------------------------------------- INTERFACE -----------------------------
 def generate_responses(query):
     # Randomize model order
     pipeline_names = list(pipelines.keys())
@@ -156,30 +430,37 @@ def reset_ui():
         gr.update(value=""),  # clear query
         gr.update(visible=False),  # hide radio
         gr.update(visible=False),  # hide vote button
-        gr.update(value=""),  # clear Option 1 text
-        gr.update(value=""),  # clear Option 2 text
+        gr.update(value="**Generating...**"),  # clear Option 1 text
+        gr.update(value="**Generating...**"),  # clear Option 2 text
         gr.update(value=""),  # clear result
         gr.update(active=False),
     )
 
 
-with gr.Blocks() as demo:
-    gr.Markdown("# Music Genre Recommendation Side-By-Side Comparison")
-    query = gr.Textbox(label="Your Query")
+app_description = load_text_resource(Path("./resources/description.md"))
+
+with gr.Blocks(title="SADAIMREC") as demo:
+    gr.Markdown(app_description)
+    query = gr.Textbox(
+        label="Your Query",
+        placeholder="Calming, music for deep relaxation with echoing sounds and deep bass",
+    )
     submit_btn = gr.Button("Submit")
     # timer that resets ui after feedback is sent
     reset_timer = gr.Timer(value=2.0, active=False)
 
     # Hidden components to store model responses and names
     with gr.Row(visible=False) as response_row:
-        response_1 = gr.Textbox(label="Option 1", interactive=False)
-        response_2 = gr.Textbox(label="Option 2", interactive=False)
+        response_1 = gr.Markdown(value="**Generating...**", label="Option 1")
+        response_2 = gr.Markdown(value="**Generating...**", label="Option 2")
     model_label_1 = gr.Textbox(visible=False)
     model_label_2 = gr.Textbox(visible=False)
 
     # Feedback
     vote = gr.Radio(
-        ["Option 1", "Option 2"], label="Select Best Response", visible=False
+        ["Option 1 (left)", "Option 2 (right)"],
+        label="Select Best Response",
+        visible=False,
     )
     vote_btn = gr.Button("Vote", visible=False)
     result = gr.Textbox(label="Console", interactive=False)
@@ -189,6 +470,7 @@ with gr.Blocks() as demo:
         fn=generate_responses,
         inputs=[query],
         outputs=[response_1, response_2, model_label_1, model_label_2],
+        show_progress="full",
     )
     submit_btn.click(  # update ui
         fn=lambda: (
@@ -222,6 +504,7 @@ with gr.Blocks() as demo:
         trigger_mode="once",
     )
 
+
 if __name__ == "__main__":
     demo.queue(max_size=10, default_concurrency_limit=1).launch(
         server_name="0.0.0.0", server_port=7860

ingest.py

@@ -9,11 +9,13 @@ from huggingface_hub import hf_hub_download
 from qdrant_client import QdrantClient
 from qdrant_client import models as qmodels
 
+VLLM_DTYPE = os.getenv("VLLM_DTYPE")
+
 DATA_PATH = Path(os.getenv("DATA_PATH"))
 DB_PATH = DATA_PATH / "db"
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-RECREATE_DB = bool(os.getenv("RECREATE_DB", "False").lower == "true")
+RECREATE_DB = bool(os.getenv("RECREATE_DB", "False").lower() == "true")
 DATA_REPO = os.getenv("DATA_REPO")
 DATA_FILENAME = os.getenv("DATA_FILENAME")
 
@@ -24,7 +26,9 @@ dense_batch_size = 128
 sparse_batch_size = 256
 
 dense_encoder = SentenceTransformer(
-    model_name_or_path="mixedbread-ai/mxbai-embed-large-v1", device="cuda"
+    model_name_or_path="mixedbread-ai/mxbai-embed-large-v1",
+    device="cuda",
+    model_kwargs={"torch_dtype": VLLM_DTYPE},
 )
 sparse_encoder = SparseTextEmbedding(model_name="Qdrant/bm25", cuda=True)
 

prompts/api.txt

@@ -1,7 +0,0 @@
-# Purpose
-
-Recommend 5 genres based on the user query
-
-# Query
-
-{query}

resources/description.md

@@ -0,0 +1,33 @@
+# Music Genre Recommendation Side-By-Side Comparison
+
+This simple application was developed and deployed as **complementary material for my thesis**.
+
+In case of any complications, questions or suggestions, please reach out via [email](mailto:[email protected]).
+
+## Instructions
+
+1. Formulate a **search query** with description of a music genre you would like to listen to. Expected format is described below.
+2. Explore **two generated recommendation rankings**: one is created by my system, one is generated using `gemini-2.0-flash`. Order is **randomized** each run.
+3. Determine which ranking you prefer more.
+4. Vote for your choice.
+5. Wait for refresh and repeat as many times as you want.
+
+## Expected Query Format
+
+- The system was designed to support **3 categories** of music genre descriptors:
+  - **Subjective**: Emotional & perceptual qualities, desired **inner feeling** (melancholic, energetic)
+  - **Purpose-Based**: Listening setting, context, suitable activities, scenario (party, workout)
+  - **Technical**: Musical & production attributes, **HOW the sound is made** (instrumentation, timbre, tempo, lo-fi)
+- **Other descriptors are out of scope of the current implementation**:
+  - I kindly ask you to only use the above 3 categories for your queries
+  - Usage of cultural, historical, etc. descriptors can lead to suboptimal results
+- You can make the descriptors **as complex and poetic as you want**, but I kindly ask you to **limit your query to a couple of sentences**
+
+## Query Examples
+
+- `Music for deep relaxation with echoing sounds and heavy bass, perfect for unwinding after along day`
+- `Music that feels like the echo of a forgotten world—slow, sorrowful. Guitars and distant vocals create the sensation of a long, drifting sleep on the edge of melancholy and oblivion.A soundtrack to isolation, it slowly pulls you into the depths of existential despair.`
+- `Raw and filled with aggression, high-energy drums, mosh-pit vibes, high bpm, guitars`
+- `Music to study to, relaxing, chill with calm drums, some piano, and suitable for background`
+- `Creamy and cozy, suitable for evenings with loved ones`
+- `Dreamy instrumental music for midnight melancholia`

resources/prompt_api.md

@@ -0,0 +1,12 @@
+# Purpose and Context
+
+Given a user-generated Search Query describing music they wish to explore, create a ranking of the most suitable music genres.
+
+# Instructions
+
+1. Create a music genre ranking, including 5 the most suitable music genres, ordered from the most to the least suitable.
+2. Respond in JSON.
+
+# Search Query
+
+{query}

prompts/local.txt → resources/prompt_vllm.md

@@ -228,4 +228,4 @@ Given a user-generated Search Query describing music they wish to explore, you m
 
 # Search Query
 
-{query}
+{query}