Add stage 2 embeddings

src/app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import time
-from typing import List, Dict, Tuple
+from typing import List, Dict, Tuple, Callable
 from pathlib import Path
 import os
 from config import (
@@ -9,7 +9,15 @@ from config import (
     EXAMPLE_QUERIES_BY_CATEGORY,
 )
 from src.search.bm25_lexical_search import search_bm25
+from src.search.vector_search import search_vector
 from src.data_prep.data_prep import load_clean_amazon_product_data
+from src.constants.code_snippets import (
+    CODE_STAGE_1,
+    CODE_STAGE_2,
+    CODE_STAGE_3,
+    CODE_STAGE_4,
+)
+
 
 _FILE_PATH = Path(__file__).parents[1]
 
@@ -59,6 +67,26 @@ def format_results(results: List[Dict], stage_name: str, metrics: Dict) -> str:
     """
     html_parts = [f"## 🔍 {stage_name}\n\n"]
 
+    # Performance metrics at the top with prominent styling
+    html_parts.append(
+        f"""
+<div style="display: flex; gap: 20px; margin-bottom: 28px;">
+    <div class="metric-box" style="flex: 1;">
+        <div style="color: #6720FF; font-size: 0.9em; font-weight: 600; margin-bottom: 6px; letter-spacing: 0.5px;">SEMANTIC MATCH</div>
+        <div style="font-size: 2.2em; font-weight: 700; color: #1E293B;">{metrics['semantic_match']:.3f}</div>
+        <div style="color: #64748B; font-size: 0.8em; margin-top: 4px;">Higher is better</div>
+    </div>
+    <div class="metric-box" style="flex: 1;">
+        <div style="color: #6720FF; font-size: 0.9em; font-weight: 600; margin-bottom: 6px; letter-spacing: 0.5px;">LATENCY</div>
+        <div style="font-size: 2.2em; font-weight: 700; color: #1E293B;">{metrics['latency_ms']}<span style="font-size: 0.45em; color: #64748B; font-weight: 400;">ms</span></div>
+        <div style="color: #64748B; font-size: 0.8em; margin-top: 4px;">Response time</div>
+    </div>
+</div>
+"""
+    )
+
+    # Results section
+    html_parts.append('<div style="margin-top: 20px;">\n\n')
     for idx, result in enumerate(results, 1):
         category = f"{result.get('main_category', 'N/A')} > {result.get('secondary_category', 'N/A')}"
         html_parts.append(
@@ -66,42 +94,30 @@ def format_results(results: List[Dict], stage_name: str, metrics: Dict) -> str:
 <div class="result-card">
 <strong>{idx}. {result['product_name']}</strong><br/>
 <span style="color: #64748B; font-size: 0.9em;">{result['description'][:150]}...</span><br/>
-<span style="color: #94A3B8; font-size: 0.85em;">Category: {category}</span><br/>
-<span style="color: #6720FF; font-weight: 600;">Score: {result['score']:.3f}</span>
+<span style="color: #94A3B8; font-size: 0.85em;">Category: {category}</span>
 </div>
 """
         )
-
-    html_parts.append("\n---\n\n### Performance Metrics\n\n")
-    html_parts.append(
-        f"""
-| Metric | Score |
-|--------|-------|
-| **Semantic Match** | {metrics['semantic_match']:.3f} |
-| **Diversity** | {metrics['diversity']:.3f} |
-| **Latency** | {metrics['latency_ms']}ms |
-"""
-    )
+    html_parts.append('</div>')
 
     return "".join(html_parts)
 
+def run_search_function_and_time(query: str, func: Callable):
+    start = time.time()
+    results = func(query)
+    latency = int((time.time() - start) * 1000)
+    return results, latency
+
 
 def search_stage_1(query: str) -> Tuple[str, Dict]:
     """Stage 1: Baseline BM25 keyword search."""
-    start_time = time.time()
-
-    results = search_bm25(query, top_k=5)
-    latency = int((time.time() - start_time) * 1000)
-
-    unique_categories = len(set(r["main_category"] for r in results)) if results else 0
-    diversity = min(1.0, unique_categories / 5.0)
+    results, latency = run_search_function_and_time(query, search_bm25)
 
     avg_score = sum(r["score"] for r in results) / len(results) if results else 0
     semantic_match = min(1.0, avg_score / 10.0)
 
     metrics = {
         "semantic_match": semantic_match,
-        "diversity": diversity,
         "latency_ms": latency,
     }
     print(f"Searched BM25 for {query} in {latency}ms")
@@ -110,30 +126,20 @@ def search_stage_1(query: str) -> Tuple[str, Dict]:
 
 
 def search_stage_2(query: str) -> Tuple[str, Dict]:
-    """Stage 2: BM25 + Vector Embeddings."""
-    start_time = time.time()
-
-    # Placeholder: Simulated embedding search with correct format
-    results = [
-        {
-            "product_name": product["title"],
-            "description": product["description"],
-            "main_category": product["category"],
-            "secondary_category": "Placeholder",
-            "score": 0.72 + (idx * 0.04),
-        }
-        for idx, product in enumerate(SAMPLE_PRODUCTS[:4])
-    ]
+    """Stage 2: Vector Embeddings using FAISS."""
+    results, latency = run_search_function_and_time(query, search_vector)
 
-    latency = int((time.time() - start_time) * 1000)
+    avg_score = sum(r["score"] for r in results) / len(results) if results else 0
+    semantic_match = avg_score
 
     metrics = {
-        "semantic_match": 0.72,
-        "diversity": 0.70,
-        "latency_ms": max(100, latency),
+        "semantic_match": semantic_match,
+        "latency_ms": latency,
     }
+    print(f"Searched vector embeddings for '{query}' in {latency}ms")
 
-    return format_results(results, "Stage 2: + Vector Embeddings", metrics), metrics
+    # Return top 5 for display
+    return format_results(results[:5], "Stage 2: Vector Embeddings", metrics), metrics
 
 
 def search_stage_3(query: str) -> Tuple[str, Dict]:
@@ -156,7 +162,6 @@ def search_stage_3(query: str) -> Tuple[str, Dict]:
 
     metrics = {
         "semantic_match": 0.81,
-        "diversity": 0.75,
         "latency_ms": max(150, latency),
     }
 
@@ -183,7 +188,6 @@ def search_stage_4(query: str) -> Tuple[str, Dict]:
 
     metrics = {
         "semantic_match": 0.88,
-        "diversity": 0.80,
         "latency_ms": max(200, latency),
     }
 
@@ -209,19 +213,19 @@ def search_all_stages(query: str) -> Tuple[str, str, str, str, str]:
 def generate_comparison_table(all_metrics: List[Dict]) -> str:
     """Generate comparison table for all stages."""
     stage_names = [
-        "Stage 1: BM25",
-        "Stage 2: + Embeddings",
-        "Stage 3: + Query Expansion",
-        "Stage 4: + Reranking",
+        "Baseline: BM25",
+        "Stage 1: + Embeddings",
+        "Stage 2: + Query Expansion",
+        "Stage 3: + Reranking",
     ]
 
     # Build markdown table
     html = "## Stage-by-Stage Comparison\n\n"
-    html += "| Stage | Semantic Match | Diversity | Latency (ms) |\n"
-    html += "|-------|----------------|-----------|---------------|\n"
+    html += "| Stage | Semantic Match | Latency (ms) |\n"
+    html += "|-------|----------------|--------------|\n"
 
     for name, metrics in zip(stage_names, all_metrics):
-        html += f"| **{name}** | {metrics['semantic_match']:.3f} | {metrics['diversity']:.3f} | {metrics['latency_ms']} |\n"
+        html += f"| **{name}** | {metrics['semantic_match']:.3f} | {metrics['latency_ms']} |\n"
 
     # Calculate improvements
     semantic_improvement = (
@@ -233,22 +237,13 @@ def generate_comparison_table(all_metrics: List[Dict]) -> str:
         if all_metrics[0]["semantic_match"] > 0
         else 0
     )
-    diversity_improvement = (
-        (
-            (all_metrics[3]["diversity"] - all_metrics[0]["diversity"])
-            / all_metrics[0]["diversity"]
-            * 100
-        )
-        if all_metrics[0]["diversity"] > 0
-        else 0
-    )
 
     html += "\n---\n\n"
     html += "## Key Insights\n\n"
-    html += f"- **Semantic Match** improves by **{semantic_improvement:.0f}%** from Stage 1 to Stage 4\n"
-    html += f"- **Diversity** increases by **{diversity_improvement:.0f}%** showing more varied results\n"
+    html += f"- **Semantic Match** improves by **{semantic_improvement:.0f}%** from baseline to final stage\n"
     html += f"- **Latency** stays under **{max(m['latency_ms'] for m in all_metrics)}ms** maintaining fast performance\n"
-    html += "- Each stage adds incremental value to search quality\n"
+    html += "- Each stage progressively enhances search relevance while keeping response times low\n"
+    html += "- Vector embeddings provide the biggest jump in semantic understanding\n"
 
     return html
 
@@ -340,106 +335,6 @@ def generate_sample_data_table() -> str:
     return html
 
 
-# Code snippets for each stage
-CODE_STAGE_1 = """
-```python
-import bm25s
-import pandas as pd
-
-# Step 1: Create BM25 index (one-time setup)
-df = pd.read_parquet("data/amazon_products.parquet")
-corpus = df["FullText"].tolist()
-corpus_tokens = bm25s.tokenize(corpus, stopwords="en")
-
-retriever = bm25s.BM25()
-retriever.index(corpus_tokens)
-retriever.save("data/bm25_index")
-
-# Step 2: Load index and search
-bm25_index = bm25s.BM25.load("data/bm25_index", load_corpus=False)
-query_tokens = bm25s.tokenize(query, stopwords="en")
-results, scores = bm25_index.retrieve(query_tokens, k=5)
-
-# Extract top results
-top_products = [df.iloc[idx] for idx in results[0]]
-```
-"""
-
-CODE_STAGE_2 = """
-```python
-from openai import OpenAI
-import faiss
-import numpy as np
-
-client = OpenAI(
-    base_url="https://api.fireworks.ai/inference/v1"
-)
-
-# Generate embeddings
-response = client.embeddings.create(
-    model="accounts/fireworks/models/qwen3-embedding-8b",
-    input=[query] + documents
-)
-
-# Extract embeddings
-query_emb = np.array(response.data[0].embedding)
-doc_embs = np.array([d.embedding for d in response.data[1:]])
-
-# FAISS search
-index = faiss.IndexFlatIP(doc_embs.shape[1])
-index.add(doc_embs)
-scores, indices = index.search(query_emb.reshape(1, -1), k=5)
-```
-"""
-
-CODE_STAGE_3 = """
-```python
-# Query expansion with LLM
-response = client.chat.completions.create(
-    model="accounts/fireworks/models/llama-v3p1-8b-instruct",
-    messages=[{
-        "role": "user",
-        "content": f"Extract 2-3 key search concepts from: {query}"
-    }]
-)
-
-expanded_query = response.choices[0].message.content
-
-# Search with expanded query
-response = client.embeddings.create(
-    model="accounts/fireworks/models/qwen3-embedding-8b",
-    input=[expanded_query] + documents
-)
-
-# Continue with embedding search...
-```
-"""
-
-CODE_STAGE_4 = """
-```python
-# First get top 20 candidates from Stage 3
-top_20_results = get_stage_3_results(query, k=20)
-
-# Rerank with Fireworks reranker
-rerank_response = client.post(
-    "https://api.fireworks.ai/inference/v1/rerank",
-    json={
-        "model": "fireworks/qwen3-reranker-8b",
-        "query": query,
-        "documents": [r["text"] for r in top_20_results],
-        "top_n": 5
-    }
-)
-
-# Get final ranked results
-final_results = [
-    top_20_results[r["index"]]
-    for r in rerank_response.json()["results"]
-]
-```
-"""
-
-
 # Build Gradio Interface
 with gr.Blocks(
     css=CUSTOM_CSS, theme=GRADIO_THEME, title="Search Alchemy - Fireworks AI"

src/constants/code_snippets.py

@@ -0,0 +1,120 @@
+"""
+Code snippets for displaying implementation examples in the Gradio UI.
+Each snippet shows the actual implementation approach for each search stage.
+"""
+
+CODE_STAGE_1 = """
+```python
+import bm25s
+import pandas as pd
+
+# Step 1: Create BM25 index (one-time setup)
+df = pd.read_parquet("data/amazon_products.parquet")
+corpus = df["FullText"].tolist()
+corpus_tokens = bm25s.tokenize(corpus, stopwords="en")
+
+retriever = bm25s.BM25()
+retriever.index(corpus_tokens)
+retriever.save("data/bm25_index")
+
+# Step 2: Load index and search
+bm25_index = bm25s.BM25.load("data/bm25_index", load_corpus=False)
+query_tokens = bm25s.tokenize(query, stopwords="en")
+results, scores = bm25_index.retrieve(query_tokens, k=5)
+
+# Extract top results
+top_products = [df.iloc[idx] for idx in results[0]]
+```
+"""
+
+CODE_STAGE_2 = """
+```python
+from openai import OpenAI
+import faiss
+import numpy as np
+
+# Initialize Fireworks AI client
+client = OpenAI(
+    api_key="your_fireworks_api_key",
+    base_url="https://api.fireworks.ai/inference/v1"
+)
+
+# Generate query embedding
+response = client.embeddings.create(
+    model="accounts/fireworks/models/qwen3-embedding-8b",
+    input=query
+)
+query_embedding = np.array(response.data[0].embedding, dtype=np.float32)
+query_vector = query_embedding.reshape(1, -1)
+
+# Normalize for cosine similarity using L2 distance
+faiss.normalize_L2(query_vector)
+
+# Load pre-built FAISS index
+index = faiss.read_index("data/faiss_index.bin")
+
+# Search for top-k similar documents
+distances, indices = index.search(query_vector, k=10)
+
+# Convert L2 distances to cosine similarity scores
+# After normalization: L2_distance = 2 * (1 - cosine_similarity)
+# So: cosine_similarity = 1 - (L2_distance / 2)
+similarity_scores = 1 - (distances[0] / 2)
+
+# Get top results
+top_results = [
+    {
+        "product": df.iloc[idx],
+        "score": float(score)
+    }
+    for idx, score in zip(indices[0], similarity_scores)
+]
+```
+"""
+
+CODE_STAGE_3 = """
+```python
+# Query expansion with LLM
+response = client.chat.completions.create(
+    model="accounts/fireworks/models/llama-v3p1-8b-instruct",
+    messages=[{
+        "role": "user",
+        "content": f"Extract 2-3 key search concepts from: {query}"
+    }]
+)
+
+expanded_query = response.choices[0].message.content
+
+# Search with expanded query
+response = client.embeddings.create(
+    model="accounts/fireworks/models/qwen3-embedding-8b",
+    input=[expanded_query] + documents
+)
+
+# Continue with embedding search...
+```
+"""
+
+CODE_STAGE_4 = """
+```python
+# First get top 20 candidates from Stage 3
+top_20_results = get_stage_3_results(query, k=20)
+
+# Rerank with Fireworks reranker
+rerank_response = client.post(
+    "https://api.fireworks.ai/inference/v1/rerank",
+    json={
+        "model": "fireworks/qwen3-reranker-8b",
+        "query": query,
+        "documents": [r["text"] for r in top_20_results],
+        "top_n": 5
+    }
+)
+
+# Get final ranked results
+final_results = [
+    top_20_results[r["index"]]
+    for r in rerank_response.json()["results"]
+]
+```
+"""

src/search/bm25_lexical_search.py

@@ -1,7 +1,7 @@
 import bm25s
 from typing import List, Dict
 from pathlib import Path
-from src.data_prep.constants import BM25_INDEX, PRODUCTS_DF
+from constants.constants import BM25_INDEX, PRODUCTS_DF
 
 _FILE_PATH = Path(__file__).parents[2]
 

src/search/vector_search.py

@@ -0,0 +1,51 @@
+import numpy as np
+import faiss
+from typing import List, Dict
+from pathlib import Path
+from src.fireworks.inference import get_embedding
+from constants.constants import FAISS_INDEX, PRODUCTS_DF
+
+_FILE_PATH = Path(__file__).parents[2]
+
+
+def search_vector(query: str, top_k: int = 10) -> List[Dict[str, any]]:
+    """
+    Search products using vector embeddings and FAISS for semantic search.
+
+    This is Stage 2: semantic search using vector embeddings to understand
+    query meaning and intent beyond exact keyword matching.
+
+    Args:
+        query: Search query string
+        top_k: Number of top results to return (default: 10)
+
+    Returns:
+        List of dictionaries containing product information and scores
+    """
+    query_embedding = get_embedding(query)
+    query_vector = np.array([query_embedding], dtype=np.float32)
+
+    # Normalize query vector for cosine similarity
+    faiss.normalize_L2(query_vector)
+
+    # Unpack FAISS index tuple (index, embeddings)
+    faiss_index = FAISS_INDEX[0]
+
+    # Search FAISS index
+    distances, indices = faiss_index.search(query_vector, top_k)
+
+    # Convert L2 distances to similarity scores (0-1 range)
+    # After normalization, L2 distance = 2 * (1 - cosine_similarity)
+    # So cosine_similarity = 1 - (L2_distance / 2)
+    similarity_scores = 1 - (distances[0] / 2)
+
+    return [
+        {
+            "product_name": PRODUCTS_DF.iloc[idx]["Product Name"],
+            "description": PRODUCTS_DF.iloc[idx]["Description"],
+            "main_category": PRODUCTS_DF.iloc[idx]["MainCategory"],
+            "secondary_category": PRODUCTS_DF.iloc[idx]["SecondaryCategory"],
+            "score": float(score),
+        }
+        for idx, score in zip(indices[0], similarity_scores)
+    ]