Update app.py

app.py

@@ -122,7 +122,7 @@ def retrieve_gradio(knowledge_base: str, query: str, topk: int):
         for line in f:
             md5s.append(line.rstrip('\n'))
     
-    doc_reps = np.load(os.path.join(target_cache_dir, f"reps.npy"))
+    doc_reps = np.load(os.path.join(target_cache_dir, f"reps.npy")).squeeze(1)
 
     query_with_instruction = "Represent this query for retrieving relavant document: " + query
     with torch.no_grad():