Update app.py

app.py

@@ -1,82 +1,100 @@
 import gradio as gr
-from sentence_transformers import CrossEncoder
+from sentence_transformers import SentenceTransformer, CrossEncoder, util
 import torch
 import requests
+import ast
 
 # -------------------------------
-# CONFIG
+# MODELS
 # -------------------------------
-HF_MODEL = "cross-encoder/ms-marco-MiniLM-L-12-v2"
+BI_ENCODER = "sentence-transformers/all-MiniLM-L6-v2"
+CROSS_ENCODER_RERANK = "cross-encoder/ms-marco-MiniLM-L-12-v2"
+CROSS_ENCODER_STS = "cross-encoder/stsb-roberta-large"
+CROSS_ENCODER_NLI = "cross-encoder/nli-deberta-v3-base"
 JINA_MODEL = "jina-reranker-m0"
 JINA_API_KEY = "jina_4075150fa702471c85ddea0a9ad4b306ouE7ymhrCpvxTxX3mScUv5LLDPKQ"
 JINA_ENDPOINT = "https://api.jina.ai/v1/rerank"
 
 # -------------------------------
-# Load Hugging Face CrossEncoder
+# Load models
 # -------------------------------
-hf_model = CrossEncoder(HF_MODEL)
+bi_encoder = SentenceTransformer(BI_ENCODER)
+ce_rerank = CrossEncoder(CROSS_ENCODER_RERANK)
+ce_sts = CrossEncoder(CROSS_ENCODER_STS)
+ce_nli = CrossEncoder(CROSS_ENCODER_NLI, num_labels=3)
 
-def rerank(query, docs_text):
-    # Split input documents (one per line)
-    docs = [d.strip() for d in docs_text.split("\n") if d.strip()]
-    if not docs:
-        return "⚠️ No documents provided."
+# -------------------------------
+# Pipeline Function
+# -------------------------------
+def evaluate_models(query, docs_str):
+    try:
+        # Parse docs string as Python list
+        docs = ast.literal_eval(docs_str)
+        assert isinstance(docs, list), "Input must be a Python list of strings"
+    except Exception as e:
+        return f"⚠️ Error parsing documents list: {e}"
 
-    # -------------------------------
-    # Hugging Face CrossEncoder Scores
-    # -------------------------------
-    hf_scores = hf_model.predict([(query, d) for d in docs])
-    hf_scores = [torch.sigmoid(torch.tensor(s)).item() for s in hf_scores]
-    hf_ranking = sorted(zip(docs, hf_scores), key=lambda x: x[1], reverse=True)
+    results = {}
 
-    # -------------------------------
-    # Jina Reranker API Scores
-    # -------------------------------
-    headers = {
-        "Authorization": f"Bearer {JINA_API_KEY}",
-        "Content-Type": "application/json",
-    }
-    payload = {
-        "model": JINA_MODEL,
-        "query": query,
-        "documents": docs,
-    }
+    # 1. Bi-encoder cosine similarity
+    query_emb = bi_encoder.encode(query, convert_to_tensor=True)
+    doc_embs = bi_encoder.encode(docs, convert_to_tensor=True)
+    cos_scores = util.cos_sim(query_emb, doc_embs)[0].cpu().tolist()
+    results["1. Bi-encoder similarity"] = sorted(zip(docs, cos_scores), key=lambda x: x[1], reverse=True)
+
+    # 2. CrossEncoder reranker (MS MARCO)
+    ce_rerank_scores = ce_rerank.predict([(query, d) for d in docs])
+    ce_rerank_scores = [torch.sigmoid(torch.tensor(s)).item() for s in ce_rerank_scores]
+    results["2. CrossEncoder Reranker (MS MARCO)"] = sorted(zip(docs, ce_rerank_scores), key=lambda x: x[1], reverse=True)
+
+    # 3. Jina Reranker
+    headers = {"Authorization": f"Bearer {JINA_API_KEY}", "Content-Type": "application/json"}
+    payload = {"model": JINA_MODEL, "query": query, "documents": docs}
     try:
-        r = requests.post(JINA_ENDPOINT, headers=headers, json=payload, timeout=20)
+        r = requests.post(JINA_ENDPOINT, headers=headers, json=payload, timeout=30)
         r.raise_for_status()
-        results = r.json()["results"]
-        jina_scores = [res["relevance_score"] for res in results]
-        jina_ranking = sorted(zip(docs, jina_scores), key=lambda x: x[1], reverse=True)
+        jina_scores = [res["relevance_score"] for res in r.json()["results"]]
+        results["3. Jina Reranker"] = sorted(zip(docs, jina_scores), key=lambda x: x[1], reverse=True)
     except Exception as e:
-        jina_ranking = [("Error", str(e))]
+        results["3. Jina Reranker"] = [(f"Error: {e}", 0)]
+
+    # 4. CrossEncoder STS
+    ce_sts_scores = ce_sts.predict([(query, d) for d in docs])
+    results["4. CrossEncoder STS"] = sorted(zip(docs, ce_sts_scores), key=lambda x: x[1], reverse=True)
+
+    # 5. CrossEncoder NLI
+    ce_nli_probs = ce_nli.predict([(query, d) for d in docs], apply_softmax=True)
+    ce_nli_scores = [float(p[1] + p[2]) for p in ce_nli_probs]  # neutral + entailment
+    results["5. CrossEncoder NLI"] = sorted(zip(docs, ce_nli_scores), key=lambda x: x[1], reverse=True)
+
+    # 6. Bi-encoder raw similarity (duplicate for clarity)
+    results["6. Bi-encoder baseline"] = sorted(zip(docs, cos_scores), key=lambda x: x[1], reverse=True)
 
     # -------------------------------
     # Format output
     # -------------------------------
-    out = "### Hugging Face Ranking\n"
-    for doc, score in hf_ranking:
-        out += f"- ({score:.4f}) {doc}\n"
-
-    out += "\n### Jina Reranker Ranking\n"
-    for doc, score in jina_ranking:
-        out += f"- ({score}) {doc}\n"
-
+    out = ""
+    for model_name, ranked in results.items():
+        out += f"\n### {model_name}\n"
+        for doc, score in ranked:
+            out += f"- ({round(score,4)}) {doc}\n"
     return out
 
 # -------------------------------
-# Simple UI
+# Gradio UI
 # -------------------------------
-with gr.Blocks() as demo:
-    gr.Markdown("### 🔎 Query + Multiple Docs Reranking (HF vs Jina)")
-    query = gr.Textbox(label="Query", lines=2, placeholder="Enter your query here...")
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("## 🔎 Multi-Model Reranker (HF + Jina)\nPass a **query** and a **list of documents (Python list of strings)**.")
+
+    query = gr.Textbox(label="Query", lines=2, placeholder="Enter your search query...")
     docs = gr.Textbox(
-        label="Candidate Documents (one per line)", 
-        lines=10, 
-        placeholder="Paste multiple document chunks here, each on a new line..."
+        label="Documents (Python list)", 
+        lines=6, 
+        placeholder='Example: ["Doc one text", "Doc two text", "Doc three text"]'
     )
-    out = gr.Textbox(label="Ranked Results", lines=15)
+    out = gr.Textbox(label="Ranked Results", lines=20)
 
-    btn = gr.Button("Rerank 🚀")
-    btn.click(rerank, inputs=[query, docs], outputs=out)
+    btn = gr.Button("Evaluate 🚀")
+    btn.click(evaluate_models, inputs=[query, docs], outputs=out)
 
 demo.launch()