Update app.py

app.py

@@ -1,23 +1,68 @@
 import gradio as gr
 from sentence_transformers import SentenceTransformer
+from rank_bm25 import BM25Okapi
+from transformers import AutoTokenizer, AutoModel
+import torch
 
-# Load the multilingual embedding model
-model = SentenceTransformer('sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2')
+# 1. Dense embedding model (HF bi-encoder)
+dense_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
 
-# Define a function to embed text
-def embed(text: str):
+def embed_dense(text: str):
     if not text.strip():
         return {"error": "Input text is empty."}
-    embedding = model.encode([text])[0]  # Get the embedding vector
-    return {"embedding": embedding.tolist()}
-
-# Launch Gradio interface
-demo = gr.Interface(
-    fn=embed,
-    inputs=gr.Textbox(lines=3, label="Input Text"),
-    outputs="json",
-    title="Multilingual Text Embedder",
-    description="Uses paraphrase-multilingual-MiniLM-L12-v2 to convert text into embeddings"
-)
-
-demo.launch()
+    emb = dense_model.encode([text])[0]
+    return {"dense_embedding": emb.tolist()}
+
+# 2. Sparse embedding model (BM25)
+# Uses rank_bm25 to compute term weights
+
+def embed_sparse(text: str):
+    if not text.strip():
+        return {"error": "Input text is empty."}
+    tokens = text.split()
+    bm25 = BM25Okapi([tokens])
+    scores = bm25.get_scores(tokens)
+    # Map each term to its BM25 weight
+    term_weights = {tok: float(score) for tok, score in zip(tokens, scores)}
+    return {"sparse_embedding": term_weights}
+
+# 3. Late-interaction embedding model (ColBERT)
+colbert_tokenizer = AutoTokenizer.from_pretrained('colbert-ir/colbertv2.0', use_fast=True)
+colbert_model = AutoModel.from_pretrained('colbert-ir/colbertv2.0')
+
+# Freeze model parameters for inference speed
+for param in colbert_model.parameters():
+    param.requires_grad = False
+
+
+def embed_colbert(text: str):
+    if not text.strip():
+        return {"error": "Input text is empty."}
+    inputs = colbert_tokenizer(text, return_tensors='pt', truncation=True, max_length=64)
+    with torch.no_grad():
+        outputs = colbert_model(**inputs)
+    # last_hidden_state: (1, seq_len, hidden_size)
+    embeddings = outputs.last_hidden_state.squeeze(0).tolist()
+    return {"colbert_embeddings": embeddings}
+
+# Build Gradio interface with tabs for each model
+with gr.Blocks(title="Text Embedding Playground") as demo:
+    gr.Markdown("# Text Embedding Playground\nChoose a model and input text to get embeddings.")
+    with gr.Tab("Dense (MiniLM-L6-v2)"):
+        txt1 = gr.Textbox(lines=3, label="Input Text")
+        out1 = gr.JSON(label="Embedding")
+        txt1.submit(embed_dense, txt1, out1)
+        gr.Button("Embed").click(embed_dense, txt1, out1)
+    with gr.Tab("Sparse (BM25)"):
+        txt2 = gr.Textbox(lines=3, label="Input Text")
+        out2 = gr.JSON(label="Term Weights")
+        txt2.submit(embed_sparse, txt2, out2)
+        gr.Button("Embed").click(embed_sparse, txt2, out2)
+    with gr.Tab("Late-Interaction (ColBERT)"):
+        txt3 = gr.Textbox(lines=3, label="Input Text")
+        out3 = gr.JSON(label="Embeddings per Token")
+        txt3.submit(embed_colbert, txt3, out3)
+        gr.Button("Embed").click(embed_colbert, txt3, out3)
+
+if __name__ == "__main__":
+    demo.launch()