Update app.py

app.py

@@ -88,7 +88,7 @@ def add_pdf_gradio(pdf_file_binary, progress=gr.Progress()):
         image_md5 = get_image_md5(image)
         image_md5s.append(image_md5)
         with torch.no_grad():
-            reps = encode([image]).squeeze(0)
+            reps = encode([image])
         reps_list.append(reps)
         images.append(image)
 
@@ -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")).squeeze(1)
+    doc_reps = np.load(os.path.join(target_cache_dir, f"reps.npy"))
 
     query_with_instruction = "Represent this query for retrieving relavant document: " + query
     with torch.no_grad():