Update app.py

app.py

@@ -2,11 +2,11 @@ import gradio as gr
 from sentence_transformers import SentenceTransformer
 import torch
 
-# Cache loaded models to avoid reloading
+# Cache loaded models
 loaded_models = {}
 
 def load_model(model_name):
-    """Load and cache a model."""
+    """Load and cache a SentenceTransformer model."""
     if model_name in loaded_models:
         return loaded_models[model_name]
     model = SentenceTransformer(model_name)
@@ -20,16 +20,16 @@ def find_similar_documents(query, documents, model_name):
         return "⚠️ Please enter a query."
     if not documents.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 valid document."
-    
+
     query_emb = model.encode_query(query)
     doc_emb = model.encode_document(doc_list)
     similarities = model.similarity(query_emb, doc_emb)
-    
+
     sorted_idx = torch.argsort(similarities[0], descending=True)
     results = []
     for i, idx in enumerate(sorted_idx):
@@ -44,7 +44,7 @@ def compare_models(query, documents, tarka_model, open_model):
         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)
 
@@ -52,7 +52,6 @@ def compare_models(query, documents, tarka_model, open_model):
     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)
@@ -61,31 +60,40 @@ def compare_models(query, documents, tarka_model, open_model):
     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]}")
+    # Make them look like cards
+    def format_result(sorted_indices, sims, model_label):
+        res = [
+            f"<div style='background-color:#f9fafb;border-radius:12px;padding:10px 14px;margin-bottom:10px;border:1px solid #e5e7eb;'>"
+            f"<b>{i+1}. (Score: {sims[0][idx]:.4f})</b><br>{doc_list[idx]}"
+            f"</div>"
+            for i, idx in enumerate(sorted_indices)
+        ]
+        return f"<div style='font-family:Inter,sans-serif;font-size:15px;line-height:1.5;'>{''.join(res)}</div>"
 
-    for i, idx in enumerate(osorted):
-        open_results.append(f"**{i+1}. (Score: {osim[0][idx]:.4f})**\n{doc_list[idx]}")
+    tarka_html = format_result(tsorted, tsim, "Tarka Model")
+    open_html = format_result(osorted, osim, "Open Model")
 
-    return "\n\n".join(tarka_results), "\n\n".join(open_results)
+    return tarka_html, open_html
 
 
-# --- UI Layout ---
+# --------------------------- UI ---------------------------
 with gr.Blocks(
     title="Document Similarity Explorer",
-    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="indigo", neutral_hue="zinc")
+    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="indigo", neutral_hue="zinc", font=[gr.themes.GoogleFont("Inter"), "Inter", "sans-serif"]),
 ) as demo:
     
-    gr.Markdown("# 🔍 Document Similarity Explorer")
-    gr.Markdown("Compare document relevance across embedding models easily.")
-    
+    gr.Markdown(
+        """
+        # 🧠 Tarka Embedding Model Playground
+        Experiment with Tarka-AIR’s embedding family for semantic search and compare performance with open-source baselines.
+        """,
+    )
+
     with gr.Tabs():
-        # ----------------- SINGLE MODEL TAB -----------------
+        # ---------------- SINGLE MODEL SEARCH ----------------
         with gr.Tab("Single Model Search"):
             with gr.Row():
                 with gr.Column(scale=1):
@@ -100,21 +108,21 @@ with gr.Blocks(
                         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")
 
@@ -122,7 +130,7 @@ with gr.Blocks(
                 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)
-            
+
             model_selector.change(fn=on_model_change, inputs=[model_selector], outputs=[loading_msg])
 
             search_btn.click(fn=find_similar_documents,
@@ -133,16 +141,16 @@ with gr.Blocks(
                                inputs=[query_input, docs_input, model_selector],
                                outputs=result_box)
 
-        # ----------------- COMPARISON TAB -----------------
+        # ---------------- MODEL COMPARISON ----------------
         with gr.Tab("Compare Models"):
+            gr.Markdown("### ⚖️ Compare how different models rank the same documents")
+            
             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"
                     )
@@ -173,9 +181,10 @@ with gr.Blocks(
                     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")
+                    gr.Markdown("#### 📊 Comparison Results")
+                    with gr.Row(equal_height=True):
+                        tarka_output = gr.HTML(label="Tarka Model Results")
+                        open_output = gr.HTML(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)
@@ -198,16 +207,15 @@ with gr.Blocks(
                                  inputs=[query_compare, docs_compare, tarka_selector, open_selector],
                                  outputs=[tarka_output, open_output])
 
-    # Example block for both tabs
+    # ---------------- Example Section ----------------
     gr.Examples(
         examples=[
             [
                 "Which planet is known as the Red Planet?",
-                "Venus is Earth's twin.\nMars, the Red Planet.\nJupiter is the biggest.\nSaturn has rings.",
-                "Tarka-AIR/Tarka-Embedding-150M-V1"
+                "Venus is Earth's twin.\nMars, the Red Planet.\nJupiter is the biggest.\nSaturn has rings."
             ]
         ],
-        inputs=[query_input, docs_input, model_selector],
+        inputs=[query_input, docs_input],
         label="Try Example"
     )