Create app.py

app.py

@@ -0,0 +1,187 @@
+import os, json, re
+from typing import List, Tuple
+
+import numpy as np
+import gradio as gr
+import faiss
+from sentence_transformers import SentenceTransformer
+
+# ---------- Paths (expects files committed under ./assets) ----------
+APP_DIR    = os.path.dirname(__file__)
+ASSETS_DIR = os.path.join(APP_DIR, "assets")
+CACHE_DIR  = "/mnt/data/eg_space_cache"  # runtime cache
+os.makedirs(CACHE_DIR, exist_ok=True)
+
+CORPUS_JSON = os.path.join(ASSETS_DIR, "corpus.json")
+EMB_FP32    = os.path.join(ASSETS_DIR, "doc_embs_fp32.npy")
+EMB_FP16    = os.path.join(ASSETS_DIR, "doc_embs_fp16.npy")
+FAISS_MAIN  = os.path.join(ASSETS_DIR, "faiss_ip_768.index")
+
+# ---------- Matryoshka dims ----------
+MATRYOSHKA_DIMS = [768, 512, 256, 128]
+DEFAULT_DIMS    = 768
+
+# ---------- Load corpus ----------
+with open(CORPUS_JSON, "r", encoding="utf-8") as f:
+    corpus = json.load(f)  # list of {"title","text"} in EXACT same order as embeddings
+
+# ---------- Load embeddings ----------
+if os.path.exists(EMB_FP32):
+    doc_embs = np.load(EMB_FP32).astype(np.float32, copy=False)
+elif os.path.exists(EMB_FP16):
+    doc_embs = np.load(EMB_FP16).astype(np.float32)  # cast back for FAISS
+else:
+    raise FileNotFoundError("Expected assets/doc_embs_fp32.npy or assets/doc_embs_fp16.npy")
+
+if doc_embs.ndim != 2 or doc_embs.shape[0] != len(corpus):
+    raise ValueError("Embeddings shape mismatch vs corpus length.")
+
+EMB_DIM = doc_embs.shape[1]  # should be 768
+
+# ---------- Model (for queries + sentence-level ops) ----------
+model = SentenceTransformer("google/embeddinggemma-300m")  # CPU is fine for queries
+
+# ---------- FAISS indexes ----------
+# Option A: load a ready-made 768-dim index if provided
+if os.path.exists(FAISS_MAIN):
+    base_index_768 = faiss.read_index(FAISS_MAIN)
+else:
+    base_index_768 = faiss.IndexFlatIP(EMB_DIM)
+    base_index_768.add(doc_embs.astype(np.float32, copy=False))
+
+# Build per-dimension flat IP indexes from the loaded embeddings
+class MultiDimFaiss:
+    def __init__(self, doc_embs_full: np.ndarray):
+        self.full = doc_embs_full
+        self.indexes = {}
+        for d in MATRYOSHKA_DIMS:
+            if d == 768 and FAISS_MAIN and os.path.exists(FAISS_MAIN):
+                self.indexes[d] = base_index_768
+            else:
+                view = self.full[:, :d].astype(np.float32, copy=False)
+                idx = faiss.IndexFlatIP(d)
+                idx.add(view)
+                self.indexes[d] = idx
+
+    def search(self, q_vec: np.ndarray, top_k: int, dims: int) -> Tuple[np.ndarray, np.ndarray]:
+        q = q_vec[:dims].astype(np.float32, copy=False)[None, :]
+        idx = self.indexes[dims]
+        return idx.search(q, top_k)
+
+faiss_md = MultiDimFaiss(doc_embs)
+
+# ---------- Core ops ----------
+def _format_snippet(text: str, max_len: int = 380) -> str:
+    return text[:max_len] + ("…" if len(text) > max_len else "")
+
+def do_search(query: str, top_k: int = 5, dims: int = DEFAULT_DIMS) -> List[List[str]]:
+    if not query or not query.strip():
+        return []
+    q_emb = model.encode_query(
+        query.strip(),
+        normalize_embeddings=True,
+        convert_to_numpy=True
+    )
+    scores, idxs = faiss_md.search(q_emb, top_k=top_k, dims=dims)
+    rows = []
+    for s, i in zip(scores[0].tolist(), idxs[0].tolist()):
+        if i == -1:
+            continue
+        title = corpus[i]["title"]
+        snippet = _format_snippet(corpus[i]["text"])
+        rows.append([f"{s:.4f}", title, snippet])
+    return rows
+
+def do_similarity(text_a: str, text_b: str, dims: int = DEFAULT_DIMS) -> float:
+    if not text_a or not text_b:
+        return 0.0
+    a = model.encode_document([text_a], normalize_embeddings=True, convert_to_numpy=True)[0][:dims]
+    b = model.encode_document([text_b], normalize_embeddings=True, convert_to_numpy=True)[0][:dims]
+    return float(np.dot(a, b))
+
+# Extractive summarization using EmbeddingGemma's Summarization prompt
+def _split_sents(text: str):
+    parts = re.split(r"(?<=[\.!?])\s+", text.strip())
+    return [p.strip() for p in parts if p.strip()]
+
+def summarize_extractive(text: str, n: int, dims: int, lambda_diversity: float = 0.7) -> str:
+    sents = _split_sents(text)
+    if not sents:
+        return ""
+    embs = model.encode(
+        sents,
+        prompt_name="Summarization",
+        normalize_embeddings=True,
+        convert_to_numpy=True,
+        batch_size=128,
+    )[:, :dims]
+    centroid = embs.mean(axis=0)
+    base = embs @ centroid
+
+    picked = []
+    for _ in range(min(n, len(sents))):
+        if not picked:
+            i = int(np.argmax(base))
+        else:
+            sim_to_sel = np.max(embs[picked] @ embs.T, axis=0)
+            mmr = (1 - lambda_diversity) * base + lambda_diversity * (1 - sim_to_sel)
+            i = int(np.argmax(mmr))
+        picked.append(i)
+        base[i] = -1e9
+    # keep original order
+    ordered = [s for _, s in sorted(zip(picked, [sents[i] for i in picked]))]
+    return " ".join(ordered)
+
+# ---------- Gradio UI ----------
+with gr.Blocks(title="EmbeddingGemma × Wikipedia (EN corpus)") as demo:
+    gr.Markdown(
+        "# EmbeddingGemma × Wikipedia (EN corpus)\n"
+        "Search, cross-lingual retrieval (to EN), similarity, and extractive summarization.\n"
+        "Model: `google/embeddinggemma-300m` • Docs indexed from `corpus.json` • Precomputed embeddings loaded from assets."
+    )
+
+    with gr.Tabs():
+        # 1) Semantic Search (EN-only corpus)
+        with gr.TabItem("Semantic Search (EN corpus)"):
+            with gr.Row():
+                q = gr.Textbox(label="Query", value="Who discovered penicillin?")
+                topk = gr.Slider(1, 20, value=5, step=1, label="Top-K")
+                dims = gr.Dropdown([str(d) for d in MATRYOSHKA_DIMS], value=str(DEFAULT_DIMS), label="Embedding dims")
+                run = gr.Button("Search")
+            out = gr.Dataframe(headers=["score", "title", "snippet"], wrap=True)
+            run.click(lambda query, k, d: do_search(query, int(k), int(d)), [q, topk, dims], out)
+
+        # 2) Cross-Lingual (queries in FR/ES/etc → EN corpus)
+        with gr.TabItem("Cross-Lingual (EN corpus)"):
+            gr.Markdown("Type your query in **French/Spanish/Arabic**. Results come from the **English-only** corpus.")
+            with gr.Row():
+                qx = gr.Textbox(label="Query", value="¿Quién descubrió la penicilina?")
+                topkx = gr.Slider(1, 20, value=5, step=1, label="Top-K")
+                dimsx = gr.Dropdown([str(d) for d in MATRYOSHKA_DIMS], value=str(DEFAULT_DIMS), label="Embedding dims")
+                runx = gr.Button("Search")
+            outx = gr.Dataframe(headers=["score", "title", "snippet"], wrap=True)
+            runx.click(lambda query, k, d: do_search(query, int(k), int(d)), [qx, topkx, dimsx], outx)
+
+        # 3) Similarity
+        with gr.TabItem("Similarity"):
+            with gr.Row():
+                a = gr.Textbox(lines=5, label="Text A", value="Alexander Fleming observed a mold that killed bacteria in 1928.")
+                b = gr.Textbox(lines=5, label="Text B", value="La penicilina fue descubierta por Alexander Fleming en 1928.")
+            dims2 = gr.Dropdown([str(d) for d in MATRYOSHKA_DIMS], value=str(DEFAULT_DIMS), label="Embedding dims")
+            sim_btn = gr.Button("Compute Similarity")
+            sim_out = gr.Number(label="Cosine similarity (-1..1)")
+            sim_btn.click(lambda x, y, d: do_similarity(x, y, int(d)), [a, b, dims2], sim_out)
+
+        # 4) Summarization (extractive)
+        with gr.TabItem("Summarization"):
+            gr.Markdown("**Extractive summarization** using EmbeddingGemma's `Summarization` prompt. Paste any long text.")
+            with gr.Row():
+                sum_dims = gr.Dropdown([str(d) for d in MATRYOSHKA_DIMS], value=str(DEFAULT_DIMS), label="Embedding dims")
+                sum_n = gr.Slider(1, 10, value=5, step=1, label="Sentences in summary")
+            sum_text = gr.Textbox(lines=12, label="Text to summarize", value="Paste a Wikipedia article (or any text) here…")
+            sum_btn = gr.Button("Summarize")
+            sum_out = gr.Textbox(lines=10, label="Summary")
+            sum_btn.click(lambda t, n, d: summarize_extractive(t, int(n), int(d)), [sum_text, sum_n, sum_dims], sum_out)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)