Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from sentence_transformers import SentenceTransformer, CrossEncoder, util
+from sentence_transformers import CrossEncoder
 import torch
 import requests
 import ast
@@ -7,28 +7,22 @@ import ast
 # -------------------------------
 # MODELS
 # -------------------------------
-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"
+NV_MODEL = "NV-RerankQA-Mistral-4B-v3"   # Hugging Face hosted
 
 # -------------------------------
 # Load models
 # -------------------------------
-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)
 
 # -------------------------------
 # 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:
@@ -36,39 +30,33 @@ def evaluate_models(query, docs_str):
 
     results = {}
 
-    # 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)
+    # 1. 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)
+    results["CrossEncoder (MS MARCO)"] = sorted(zip(docs, ce_rerank_scores), key=lambda x: x[1], reverse=True)
 
-    # 3. Jina Reranker
+    # 2. 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=30)
         r.raise_for_status()
         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)
+        results["Jina Reranker"] = sorted(zip(docs, jina_scores), key=lambda x: x[1], reverse=True)
     except Exception as 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)
+        results["Jina Reranker"] = [(f"Error: {e}", 0)]
 
-    # 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)
+    # 3. NV RerankQA Mistral-4B-v3 (HF Inference API)
+    try:
+        hf_endpoint = f"https://api-inference.huggingface.co/models/{NV_MODEL}"
+        headers = {"Authorization": f"Bearer YOUR_HF_API_KEY"}
+        payload = {"inputs": {"query": query, "documents": docs}}
+        r = requests.post(hf_endpoint, headers=headers, json=payload, timeout=60)
+        r.raise_for_status()
+        nv_scores = [res["score"] for res in r.json()]
+        results["NV-RerankQA-Mistral-4B-v3"] = sorted(zip(docs, nv_scores), key=lambda x: x[1], reverse=True)
+    except Exception as e:
+        results["NV-RerankQA-Mistral-4B-v3"] = [(f"Error: {e}", 0)]
 
     # -------------------------------
     # Format output
@@ -84,7 +72,7 @@ def evaluate_models(query, docs_str):
 # Gradio UI
 # -------------------------------
 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)**.")
+    gr.Markdown("## 👑 Ranking Battle (3 Models)\nCompare **NV-RerankQA-Mistral-4B-v3**, **Jina**, and **CrossEncoder**.")
 
     query = gr.Textbox(label="Query", lines=2, placeholder="Enter your search query...")
     docs = gr.Textbox(