Update app.py

app.py

@@ -32,8 +32,8 @@ def evaluate_models(query, docs_str):
 
     # 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["CrossEncoder (MS MARCO)"] = sorted(zip(docs, ce_rerank_scores), key=lambda x: x[1], reverse=True)
+    ce_rerank_scores = [round(torch.sigmoid(torch.tensor(s)).item(), 4) for s in ce_rerank_scores]
+    results["CrossEncoder (MS MARCO)"] = ce_rerank_scores
 
     # 2. Jina Reranker
     headers = {"Authorization": f"Bearer {JINA_API_KEY}", "Content-Type": "application/json"}
@@ -41,10 +41,10 @@ def evaluate_models(query, docs_str):
     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["Jina Reranker"] = sorted(zip(docs, jina_scores), key=lambda x: x[1], reverse=True)
+        jina_scores = [round(res["relevance_score"], 4) for res in r.json()["results"]]
+        results["Jina Reranker"] = jina_scores
     except Exception as e:
-        results["Jina Reranker"] = [(f"Error: {e}", 0)]
+        results["Jina Reranker"] = [f"Error: {e}"]
 
     # 3. NV RerankQA Mistral-4B-v3 (HF Inference API)
     try:
@@ -53,26 +53,18 @@ def evaluate_models(query, docs_str):
         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)
+        nv_scores = [round(res["score"], 4) for res in r.json()]
+        results["NV-RerankQA-Mistral-4B-v3"] = nv_scores
     except Exception as e:
-        results["NV-RerankQA-Mistral-4B-v3"] = [(f"Error: {e}", 0)]
+        results["NV-RerankQA-Mistral-4B-v3"] = [f"Error: {e}"]
 
-    # -------------------------------
-    # Format output
-    # -------------------------------
-    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
+    return results
 
 # -------------------------------
 # Gradio UI
 # -------------------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("## 👑 Ranking Battle (3 Models)\nCompare **NV-RerankQA-Mistral-4B-v3**, **Jina**, and **CrossEncoder**.")
+    gr.Markdown("## 👑 Ranking Battle (Scores Only)\nOutputs only **scores** from the 3 models.")
 
     query = gr.Textbox(label="Query", lines=2, placeholder="Enter your search query...")
     docs = gr.Textbox(
@@ -80,7 +72,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         lines=6, 
         placeholder='Example: ["Doc one text", "Doc two text", "Doc three text"]'
     )
-    out = gr.Textbox(label="Ranked Results", lines=20)
+    out = gr.JSON(label="Model Scores")
 
     btn = gr.Button("Evaluate 🚀")
     btn.click(evaluate_models, inputs=[query, docs], outputs=out)