Update app.py

app.py

@@ -1,125 +1,77 @@
-import os
+import faiss
 import json
-import numpy as np
 import gradio as gr
-import faiss
 from sentence_transformers import SentenceTransformer
 
-# === CONFIGURATION ===
-# Define paths for assets (corpus + embeddings + optional FAISS index)
-BASE_DIR = os.path.dirname(__file__)
-DATA_DIR = os.path.join(BASE_DIR, "assets")
-
-CORPUS_PATH = os.path.join(DATA_DIR, "corpus.json")          # Text data (title + article)
-EMB_PATH_32 = os.path.join(DATA_DIR, "embeddings_fp32.npy")  # Embeddings in float32 (preferred)
-EMB_PATH_16 = os.path.join(DATA_DIR, "embeddings_fp16.npy")  # Embeddings in float16 (smaller, needs casting)
-FAISS_INDEX_PATH = os.path.join(DATA_DIR, "faiss_main.index")# Precomputed FAISS index (optional, saves time)
-
-# Allow searching with different embedding sizes (for speed/accuracy tradeoff)
-DIM_OPTIONS = [768, 512, 256, 128]
-DEFAULT_DIM = 768
-
-# === LOAD CORPUS ===
-# Expecting corpus.json to be a list of dicts: [{"title": "...", "text": "..."}, ...]
+# ----------------------------------------------------
+# 1. Load FAISS Index (Prebuilt)
+# ----------------------------------------------------
+# This index already contains all document embeddings.
+INDEX_PATH = "assets/faiss_main.index"
+index = faiss.read_index(INDEX_PATH)
+
+# ----------------------------------------------------
+# 2. Load Text Corpus (ID → Content Mapping)
+# ----------------------------------------------------
+# corpus.json should be a list where each position aligns
+# with the embedding index used in FAISS.
+# Example: ["Text 1", "Text 2", "Text 3", ...]
+CORPUS_PATH = "assets/corpus.json"
 with open(CORPUS_PATH, "r", encoding="utf-8") as f:
-    ARTICLES = json.load(f)
-
-# === LOAD EMBEDDINGS ===
-# Choose FP32 if available; otherwise load FP16 and convert to FP32 (required by FAISS)
-if os.path.exists(EMB_PATH_32):
-    VECTORS = np.load(EMB_PATH_32).astype("float32", copy=False)
-elif os.path.exists(EMB_PATH_16):
-    VECTORS = np.load(EMB_PATH_16).astype("float32")
-else:
-    raise RuntimeError("No embeddings found in /assets")
-
-# Sanity check: Ensure number of embeddings matches number of articles
-assert VECTORS.shape[0] == len(ARTICLES), "Embedding and corpus length mismatch"
-
-EMB_SIZE = VECTORS.shape[1]  # e.g. 768 for MiniLM or Gemma
-
-# === LOAD OR INITIALIZE FAISS INDEX ===
-# FAISS allows ultra-fast vector search using approximate or exact nearest neighbor methods.
-if os.path.exists(FAISS_INDEX_PATH):
-    base_index = faiss.read_index(FAISS_INDEX_PATH)  # Load prebuilt index for speed
-else:
-    # Normalizing vectors turns inner product (dot product) into cosine similarity
-    faiss.normalize_L2(VECTORS)
-    base_index = faiss.IndexFlatIP(EMB_SIZE)  # Flat index = brute force but optimized C++
-    base_index.add(VECTORS)
-
-# === BUILD VARIANT INDEXES WITH REDUCED DIMENSIONS (optional speed boost) ===
-class DimensionalIndexBank:
+    CORPUS = json.load(f)
+
+# ----------------------------------------------------
+# 3. Load Sentence Transformer Model for Query Encoding
+# ----------------------------------------------------
+# ⚠️ IMPORTANT: Must be the SAME model that was used
+# to generate the original corpus embeddings stored in FAISS!
+EMBEDDING_MODEL_NAME = "google/embeddinggemma-300m"  # ← Change if you used another model
+model = SentenceTransformer(EMBEDDING_MODEL_NAME)
+
+# ----------------------------------------------------
+# 4. Search Function: Query → Top K Results
+# ----------------------------------------------------
+def search_faiss(query, top_k=5):
     """
-    Manages multiple FAISS indexes at different embedding dimensions.
-    Useful for experimenting with search speed vs accuracy.
+    Takes a text query, embeds it, searches FAISS index,
+    and returns top-k most similar corpus entries.
     """
-    def __init__(self, full_matrix):
-        self.indices = {}
-        for dim in DIM_OPTIONS:
-            # If we're using full dimension and prebuilt index exists, reuse it
-            if dim == EMB_SIZE and os.path.exists(FAISS_INDEX_PATH):
-                self.indices[dim] = base_index
-            else:
-                # Slice to first 'dim' components (Matryoshka-style compression)
-                cut = full_matrix[:, :dim].astype("float32", copy=False)
-                faiss.normalize_L2(cut)
-                idx = faiss.IndexFlatIP(dim)
-                idx.add(cut)
-                self.indices[dim] = idx
+    # Encode query text into a vector
+    query_embedding = model.encode([query], convert_to_numpy=True)
 
-    def query(self, vector, top_k, dim):
-        """
-        Search the appropriate FAISS index based on embedding dimension.
-        """
-        v = vector[:dim].reshape(1, -1).astype("float32", copy=False)
-        faiss.normalize_L2(v)
-        return self.indices[dim].search(v, top_k)
+    # Perform similarity search (returns distances and indices)
+    distances, indices = index.search(query_embedding, top_k)
 
-searcher = DimensionalIndexBank(VECTORS)
-
-# === LOAD SENTENCE TRANSFORMER ===
-# This is used *only for encoding user queries*. Document embeddings are precomputed.
-encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
-
-# === HELPERS ===
-def truncate(text, limit=300):
-    """Shorten text to a preview/snippet form."""
-    return text[:limit] + "…" if len(text) > limit else text
-
-def run_search(query, k, dim):
-    """Core search function used by Gradio."""
-    if not query.strip():
-        return []
-    q_vec = encoder.encode(query, normalize_embeddings=True)
-    scores, ids = searcher.query(q_vec, k, dim)
+    # Collect matching corpus entries
     results = []
-    for score, idx in zip(scores[0], ids[0]):
-        article = ARTICLES[idx]
-        results.append([f"{score:.4f}", article["title"], truncate(article["text"])])
-    return results
-
-# === BUILD GRADIO APP ===
-def build_app():
-    with gr.Blocks(title="Semantic Search (FAISS + Transformers)") as demo:
-        gr.Markdown("## 🔍 Fast Semantic Search Over Wikipedia")
+    for rank, idx in enumerate(indices[0]):
+        results.append(
+            f"#{rank+1} | Score: {distances[0][rank]:.4f}\n{CORPUS[idx]}"
+        )
 
-        query = gr.Textbox(label="Search Input", value="Who discovered penicillin?")
-        cols = gr.Row()
-        with cols:
-            k_slider = gr.Slider(1, 20, step=1, value=5, label="Results to Show")
-            dim_choice = gr.Dropdown([str(d) for d in DIM_OPTIONS], value=str(DEFAULT_DIM), label="Embedding Size")
+    return "\n\n".join(results)
 
-        btn = gr.Button("Search")
-        output = gr.Dataframe(headers=["Score", "Title", "Snippet"], wrap=True)
-
-        btn.click(lambda q, k, d: run_search(q, int(k), int(d)), [query, k_slider, dim_choice], output)
-
-    return demo
-
-app = build_app()
-
-# === MAIN ENTRY ===
+# ----------------------------------------------------
+# 5. Build Gradio Interface
+# ----------------------------------------------------
+def gradio_search(query, top_k):
+    if not query.strip():
+        return "Please enter a search query."
+    return search_faiss(query, top_k)
+
+demo = gr.Interface(
+    fn=gradio_search,
+    inputs=[
+        gr.Textbox(label="Enter your search query"),
+        gr.Slider(1, 20, value=5, step=1, label="Number of results"),
+    ],
+    outputs=gr.Textbox(label="Search Results"),
+    title="FAISS Semantic Search",
+    description="A simple search engine powered by FAISS + Sentence Transformers."
+)
+
+# ----------------------------------------------------
+# 6. Run App (for local testing or Spaces deployment)
+# ----------------------------------------------------
 if __name__ == "__main__":
-    # When hosted on Hugging Face, Gradio auto-manages this
-    app.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()