Update app.py

app.py

@@ -2,93 +2,213 @@ import gradio as gr
 from sentence_transformers import SentenceTransformer
 import torch
 
-# Load the model
-model = SentenceTransformer("Tarka-AIR/Tarka-Embedding-150M-V1")
+# Cache loaded models to avoid reloading
+loaded_models = {}
 
-def find_similar_documents(query, documents):
+def load_model(model_name):
+    """Load and cache a model."""
+    if model_name in loaded_models:
+        return loaded_models[model_name]
+    model = SentenceTransformer(model_name)
+    loaded_models[model_name] = model
+    return model
+
+
+def find_similar_documents(query, documents, model_name):
+    """Compute similarity ranking for one model."""
     if not query.strip():
-        return "Please enter a query."
+        return "⚠️ Please enter a query."
     if not documents.strip():
-        return "Please enter documents (one per line)."
-    
-    # Split documents by lines
-    doc_list = [doc.strip() for doc in documents.split('\n') if doc.strip()]
+        return "⚠️ Please enter documents (one per line)."
     
+    model = load_model(model_name)
+    doc_list = [d.strip() for d in documents.split("\n") if d.strip()]
     if not doc_list:
-        return "Please enter at least one document."
-    
-    # Encode query and documents
-    query_embeddings = model.encode_query(query)
-    document_embeddings = model.encode_document(doc_list)
-    
-    # Compute similarities
-    similarities = model.similarity(query_embeddings, document_embeddings)
+        return "⚠️ Please enter at least one valid document."
     
-    # Sort documents by similarity score
-    sorted_indices = torch.argsort(similarities[0], descending=True)
+    query_emb = model.encode_query(query)
+    doc_emb = model.encode_document(doc_list)
+    similarities = model.similarity(query_emb, doc_emb)
     
-    # Format results
+    sorted_idx = torch.argsort(similarities[0], descending=True)
     results = []
-    for i, idx in enumerate(sorted_indices):
+    for i, idx in enumerate(sorted_idx):
         score = similarities[0][idx].item()
-        doc = doc_list[idx]
-        results.append(f"{i+1}. Score: {score:.4f}\n   Document: {doc}")
-    
+        results.append(f"**{i+1}. (Score: {score:.4f})**\n{doc_list[idx]}")
     return "\n\n".join(results)
 
-with gr.Blocks(title="Document Similarity Search", theme=gr.themes.Default(primary_hue="blue", secondary_hue="indigo", neutral_hue="zinc", font=[gr.themes.GoogleFont("Roboto"), "Arial", "sans-serif"])) as demo:
-    gr.Markdown("# Document Similarity Search")
-    gr.Markdown("Find the most relevant documents for your query using Google's Embedding Gemma model.")
+
+def compare_models(query, documents, tarka_model, open_model):
+    """Compare two models side-by-side."""
+    if not query.strip():
+        return "⚠️ Please enter a query.", ""
+    if not documents.strip():
+        return "⚠️ Please enter documents (one per line).", ""
+    
+    tarka = load_model(tarka_model)
+    openm = load_model(open_model)
+
+    doc_list = [d.strip() for d in documents.split("\n") if d.strip()]
+    if not doc_list:
+        return "⚠️ Please enter at least one valid document.", ""
+
+    # Compute similarities for both models
+    tq = tarka.encode_query(query)
+    td = tarka.encode_document(doc_list)
+    tsim = tarka.similarity(tq, td)
+
+    oq = openm.encode_query(query)
+    od = openm.encode_document(doc_list)
+    osim = openm.similarity(oq, od)
+
+    # Sort for each model
+    tsorted = torch.argsort(tsim[0], descending=True)
+    osorted = torch.argsort(osim[0], descending=True)
+
+    tarka_results, open_results = [], []
+    for i, idx in enumerate(tsorted):
+        tarka_results.append(f"**{i+1}. (Score: {tsim[0][idx]:.4f})**\n{doc_list[idx]}")
+
+    for i, idx in enumerate(osorted):
+        open_results.append(f"**{i+1}. (Score: {osim[0][idx]:.4f})**\n{doc_list[idx]}")
+
+    return "\n\n".join(tarka_results), "\n\n".join(open_results)
+
+
+# --- UI Layout ---
+with gr.Blocks(
+    title="Document Similarity Explorer",
+    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="indigo", neutral_hue="zinc")
+) as demo:
     
-    with gr.Row():
-        with gr.Column(scale=1):
-            query_input = gr.Textbox(
-                label="Query",
-                placeholder="Enter your search query...",
-                lines=2
-            )
+    gr.Markdown("# 🔍 Document Similarity Explorer")
+    gr.Markdown("Compare document relevance across embedding models easily.")
+    
+    with gr.Tabs():
+        # ----------------- SINGLE MODEL TAB -----------------
+        with gr.Tab("Single Model Search"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    model_selector = gr.Dropdown(
+                        label="Select Embedding Model",
+                        choices=[
+                            "Tarka-AIR/Tarka-Embedding-150M-V1",
+                            "sentence-transformers/all-MiniLM-L6-v2",
+                            "intfloat/e5-base-v2",
+                            "BAAI/bge-small-en-v1.5"
+                        ],
+                        value="Tarka-AIR/Tarka-Embedding-150M-V1"
+                    )
+                    loading_msg = gr.Markdown(visible=False)
+                    
+                    query_input = gr.Textbox(
+                        label="Query",
+                        placeholder="Enter your search query...",
+                        lines=2
+                    )
+                    
+                    docs_input = gr.Textbox(
+                        label="Documents",
+                        placeholder="Enter one document per line...",
+                        lines=10
+                    )
+                    
+                    search_btn = gr.Button("🔎 Search", variant="primary")
+                
+                with gr.Column(scale=1):
+                    result_box = gr.Markdown(label="Results", elem_id="results-box")
+
+            def on_model_change(model_name):
+                loading_msg.update(value=f"⏳ Loading **{model_name}**...", visible=True)
+                load_model(model_name)
+                return gr.update(value=f"✅ {model_name} loaded and ready!", visible=True)
             
-            documents_input = gr.Textbox(
-                label="Documents",
-                placeholder="Enter documents (one per line)...",
-                lines=10
-            )
+            model_selector.change(fn=on_model_change, inputs=[model_selector], outputs=[loading_msg])
+
+            search_btn.click(fn=find_similar_documents,
+                             inputs=[query_input, docs_input, model_selector],
+                             outputs=result_box)
             
-            search_btn = gr.Button("Search", variant="primary")
-        
-        with gr.Column(scale=1):
-            output = gr.Textbox(
-                label="Results",
-                lines=15,
-                show_copy_button=True
-            )
-    
-    # Example usage
+            query_input.submit(fn=find_similar_documents,
+                               inputs=[query_input, docs_input, model_selector],
+                               outputs=result_box)
+
+        # ----------------- COMPARISON TAB -----------------
+        with gr.Tab("Compare Models"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    tarka_selector = gr.Dropdown(
+                        label="Tarka Model",
+                        choices=[
+                            "Tarka-AIR/Tarka-Embedding-150M-V1",
+                            "Tarka-AIR/Tarka-Embedding-200M-V1",
+                            "Tarka-AIR/Tarka-Embedding-300M-V1"
+                        ],
+                        value="Tarka-AIR/Tarka-Embedding-150M-V1"
+                    )
+
+                    open_selector = gr.Dropdown(
+                        label="Open Source Model",
+                        choices=[
+                            "sentence-transformers/all-MiniLM-L6-v2",
+                            "intfloat/e5-base-v2",
+                            "BAAI/bge-small-en-v1.5",
+                            "thenlper/gte-base"
+                        ],
+                        value="sentence-transformers/all-MiniLM-L6-v2"
+                    )
+
+                    compare_loading = gr.Markdown(visible=False)
+
+                    query_compare = gr.Textbox(
+                        label="Query",
+                        placeholder="Enter query to compare...",
+                        lines=2
+                    )
+                    docs_compare = gr.Textbox(
+                        label="Documents",
+                        placeholder="Enter documents (one per line)...",
+                        lines=10
+                    )
+                    compare_btn = gr.Button("⚖️ Compare Models", variant="primary")
+
+                with gr.Column(scale=2):
+                    with gr.Row():
+                        tarka_output = gr.Markdown(label="Tarka Model Results")
+                        open_output = gr.Markdown(label="Open Source Model Results")
+
+            def on_compare_models_load(tarka_model, open_model):
+                compare_loading.update(value=f"⏳ Loading **{tarka_model}** and **{open_model}**...", visible=True)
+                load_model(tarka_model)
+                load_model(open_model)
+                return gr.update(value=f"✅ Models ready for comparison!", visible=True)
+
+            tarka_selector.change(fn=on_compare_models_load,
+                                  inputs=[tarka_selector, open_selector],
+                                  outputs=[compare_loading])
+            open_selector.change(fn=on_compare_models_load,
+                                 inputs=[tarka_selector, open_selector],
+                                 outputs=[compare_loading])
+
+            compare_btn.click(fn=compare_models,
+                              inputs=[query_compare, docs_compare, tarka_selector, open_selector],
+                              outputs=[tarka_output, open_output])
+
+            query_compare.submit(fn=compare_models,
+                                 inputs=[query_compare, docs_compare, tarka_selector, open_selector],
+                                 outputs=[tarka_output, open_output])
+
+    # Example block for both tabs
     gr.Examples(
         examples=[
             [
                 "Which planet is known as the Red Planet?",
-                "Venus is often called Earth's twin because of its similar size and proximity.\nMars, known for its reddish appearance, is often referred to as the Red Planet.\nJupiter, the largest planet in our solar system, has a prominent red spot.\nSaturn, famous for its rings, is sometimes mistaken for the Red Planet."
-            ],
-            [
-                "What causes seasons on Earth?",
-                "The tilt of Earth's axis causes different parts of the planet to receive varying amounts of sunlight throughout the year.\nThe moon's gravitational pull affects ocean tides but not seasons.\nEarth's orbit around the sun is slightly elliptical, but this has minimal effect on seasons.\nThe rotation of Earth on its axis causes day and night cycles."
+                "Venus is Earth's twin.\nMars, the Red Planet.\nJupiter is the biggest.\nSaturn has rings.",
+                "Tarka-AIR/Tarka-Embedding-150M-V1"
             ]
         ],
-        inputs=[query_input, documents_input]
-    )
-    
-    search_btn.click(
-        fn=find_similar_documents,
-        inputs=[query_input, documents_input],
-        outputs=output
-    )
-    
-    # Allow Enter key to trigger search
-    query_input.submit(
-        fn=find_similar_documents,
-        inputs=[query_input, documents_input],
-        outputs=output
+        inputs=[query_input, docs_input, model_selector],
+        label="Try Example"
     )
 
-demo.launch()
+demo.launch()