Update app.py

app.py

@@ -126,23 +126,37 @@ def predict(text):
         return "No query provided"
     
     query_embedding = model.encode(text, convert_to_tensor=True)
-    top_idx = int(util.pytorch_cos_sim(query_embedding, embeddings)[0].argmax())
-    top_idx2 = int(util.pytorch_cos_sim(query_embedding, embeddings2)[0].argmax())
     
-#     result = f"""Question 1: {df.iloc[top_idx]["question"]}
-# Link 1: {df.iloc[top_idx]["link"]}
-
-# Question 2: {df2.iloc[top_idx2]["question"]}
-# Link 2: {df2.iloc[top_idx2]["link"]}"""
+    # Compute similarity scores
+    sim_scores1 = util.pytorch_cos_sim(query_embedding, embeddings)[0]
+    sim_scores2 = util.pytorch_cos_sim(query_embedding, embeddings2)[0]
+    
+    # Get top 3 indices
+    top3_idx1 = sim_scores1.topk(3).indices.cpu().numpy()
+    top3_idx2 = sim_scores2.topk(3).indices.cpu().numpy()
+    
+    # Prepare results
+    results = {
+        "top1": [],
+        "top2": []
+    }
+    
+    for idx in top3_idx1:
+        results["top1"].append({
+            "question": df.iloc[idx]["question"],
+            "link": df.iloc[idx]["link"],
+            "score": float(sim_scores1[idx])
+        })
     
-    # return result
+    for idx in top3_idx2:
+        results["top2"].append({
+            "question": df2.iloc[idx]["question"],
+            "link": df2.iloc[idx]["link"],
+            "score": float(sim_scores2[idx])
+        })
     
-    return {
-            "question1": df.iloc[top_idx]["question"],
-            "link1": df.iloc[top_idx]["link"],
-            "question2": df2.iloc[top_idx2]["question"],
-            "link2": df2.iloc[top_idx2]["link"]
-        }    
+    return results
+   
 
 # Match the EXACT structure of your working translation app
 title = "Search CSV"