Update app.py

app.py

@@ -142,14 +142,14 @@ def retrieve_gradio(knowledge_base: str, query: str, topk: int):
 
     similarities_np = similarities.cpu().numpy()
     print(f"topk_doc_ids_np: {topk_doc_ids_np}, topk_values_np: {topk_values_np}")
-    images_topk = [Image.open(os.path.join(target_cache_dir, f"{index2img_filename[idx]}.png")) for idx in topk_doc_ids_np]
+    images_topk = [Image.open(os.path.join(target_cache_dir, f"{index2img_filename[idx]}")) for idx in topk_doc_ids_np]
 
     with open(os.path.join(target_cache_dir, f"q-{query_md5}.json"), 'w') as f:
         f.write(json.dumps(
             {
                 "knowledge_base": knowledge_base,
                 "query": query,
-                "retrived_docs": [os.path.join(target_cache_dir, f"{index2img_filename[idx]}.png") for idx in topk_doc_ids_np]
+                "retrived_docs": [os.path.join(target_cache_dir, f"{index2img_filename[idx]}") for idx in topk_doc_ids_np]
             }, indent=4, ensure_ascii=False
         ))