Update app.py

app.py

@@ -121,6 +121,12 @@ df2 = pd.read_csv("cleaned2.csv")
 embeddings = torch.load("embeddings1.pt")
 embeddings2 = torch.load("embeddings2.pt")
 
+# Pre-extract DataFrame columns to avoid repeated iloc calls
+df_questions = df["question"].values
+df_links = df["link"].values
+df2_questions = df2["question"].values
+df2_links = df2["link"].values
+
 def predict(text):
     if not text or text.strip() == "":
         return "No query provided"
@@ -131,32 +137,37 @@ def predict(text):
     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()
+    # Get top 3 values and indices in one call
+    top3_scores1, top3_idx1 = sim_scores1.topk(3)
+    top3_scores2, top3_idx2 = sim_scores2.topk(3)
+    
+    # Convert to CPU once
+    top3_idx1_cpu = top3_idx1.cpu().numpy()
+    top3_idx2_cpu = top3_idx2.cpu().numpy()
+    top3_scores1_cpu = top3_scores1.cpu().numpy()
+    top3_scores2_cpu = top3_scores2.cpu().numpy()
     
-    # Prepare results
+    # Prepare results using pre-extracted arrays
     results = {
-        "top1": [],
-        "top2": []
+        "top1": [
+            {
+                "question": df_questions[idx],
+                "link": df_links[idx],
+                "score": float(score)
+            }
+            for idx, score in zip(top3_idx1_cpu, top3_scores1_cpu)
+        ],
+        "top2": [
+            {
+                "question": df2_questions[idx],
+                "link": df2_links[idx], 
+                "score": float(score)
+            }
+            for idx, score in zip(top3_idx2_cpu, top3_scores2_cpu)
+        ]
     }
     
-    for idx in top3_idx1:
-        results["top1"].append({
-            "question": df.iloc[idx]["question"],
-            "link": df.iloc[idx]["link"],
-            "score": float(sim_scores1[idx])
-        })
-    
-    for idx in top3_idx2:
-        results["top2"].append({
-            "question": df2.iloc[idx]["question"],
-            "link": df2.iloc[idx]["link"],
-            "score": float(sim_scores2[idx])
-        })
-    
     return results
-   
 
 # Match the EXACT structure of your working translation app
 title = "Search CSV"