Update app.py

app.py

@@ -16,12 +16,22 @@ JINA_ENDPOINT = "https://api.jina.ai/v1/rerank"
 # -------------------------------
 hf_model = CrossEncoder(HF_MODEL)
 
-def compare_models(query, doc):
-    # Hugging Face score
-    raw_score = hf_model.predict([(query, doc)])[0]
-    hf_score = torch.sigmoid(torch.tensor(raw_score)).item()
+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."
 
-    # Jina reranker score
+    # -------------------------------
+    # 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)
+
+    # -------------------------------
+    # Jina Reranker API Scores
+    # -------------------------------
     headers = {
         "Authorization": f"Bearer {JINA_API_KEY}",
         "Content-Type": "application/json",
@@ -29,27 +39,44 @@ def compare_models(query, doc):
     payload = {
         "model": JINA_MODEL,
         "query": query,
-        "documents": [doc],
+        "documents": docs,
     }
     try:
         r = requests.post(JINA_ENDPOINT, headers=headers, json=payload, timeout=20)
         r.raise_for_status()
-        jina_score = r.json()["results"][0]["relevance_score"]
+        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)
     except Exception as e:
-        jina_score = f"Error: {str(e)}"
+        jina_ranking = [("Error", str(e))]
+
+    # -------------------------------
+    # 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"
 
-    return f"Hugging Face ({HF_MODEL}): {round(hf_score,4)}\nJina ({JINA_MODEL}): {jina_score}"
+    return out
 
 # -------------------------------
-# Simple Lite UI
+# Simple UI
 # -------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("### 🔎 Query vs Document Similarity (HF vs Jina)")
-    query = gr.Textbox(label="Query", lines=3, placeholder="Paste your query here...")
-    doc = gr.Textbox(label="Document Chunk", lines=6, placeholder="Paste your document chunk here...")
-    out = gr.Textbox(label="Scores", lines=3)
+    gr.Markdown("### 🔎 Query + Multiple Docs Reranking (HF vs Jina)")
+    query = gr.Textbox(label="Query", lines=2, placeholder="Enter your query here...")
+    docs = gr.Textbox(
+        label="Candidate Documents (one per line)", 
+        lines=10, 
+        placeholder="Paste multiple document chunks here, each on a new line..."
+    )
+    out = gr.Textbox(label="Ranked Results", lines=15)
 
-    btn = gr.Button("Compute Similarity 🚀")
-    btn.click(compare_models, inputs=[query, doc], outputs=out)
+    btn = gr.Button("Rerank 🚀")
+    btn.click(rerank, inputs=[query, docs], outputs=out)
 
 demo.launch()