Update app.py

app.py

@@ -139,7 +139,7 @@ def retrieve_gradio(knowledge_base: str, query: str, topk: int):
 
     topk_values_np = topk_values.cpu().numpy()
 
-    topk_doc_ids_np = topk_doc_ids.cpu().numpy()
+    topk_doc_ids_np = topk_doc_ids.squeeze().cpu().numpy()
 
     similarities_np = similarities.cpu().numpy()
     print(f"topk_doc_ids_np: {topk_doc_ids_np}, topk_values_np: {topk_values_np}")