Upload app.py

app.py

@@ -135,6 +135,8 @@ def retrieve_gradio(knowledge_base: str, query: str, topk: int):
     doc_list = [f for f in os.listdir(target_cache_dir) if f.endswith('.npy')]
     doc_list = sorted(doc_list)
     doc_reps = [np.load(os.path.join(target_cache_dir, f)) for f in doc_list]
+    doc_reps_cat = torch.cat([torch.Tensor(i) for i in doc_reps], dim=0)
+    doc_reps_cat = torch.cat([i for i in doc_reps_cat], dim=0)
 
     query_with_instruction = "Represent this query for retrieving relevant document: " + query
     with torch.no_grad():
@@ -142,7 +144,6 @@ def retrieve_gradio(knowledge_base: str, query: str, topk: int):
 
     query_md5 = hashlib.md5(query.encode()).hexdigest()
 
-    doc_reps_cat = torch.cat([torch.Tensor(i) for i in doc_reps], dim=0)
     print(f"query_rep_shape: {query_rep.shape}, doc_reps_cat_shape: {doc_reps_cat.shape}")
     similarities = torch.matmul(query_rep, doc_reps_cat.T)