Update pipeline.py

pipeline.py

@@ -22,7 +22,7 @@ def get_image_md5(img: Image.Image):
     hex_digest = hash_md5.hexdigest()
     return hex_digest
 
-def pdf_to_images(pdf_path, dpi=100):
+def pdf_to_images(pdf_path, dpi=200):
     doc = fitz.open(pdf_path)
     images = []
     for page in tqdm.tqdm(doc):
@@ -67,7 +67,7 @@ class PDFVisualRetrieval:
         images_topk = [all_images_doc_list[idx] for idx in topk_doc_ids_np]
         return topk_doc_ids_np, topk_values_np, images_topk
     
-    def add_pdf(self, knowledge_base_name: str, pdf_file_path: str, dpi: int = 100):
+    def add_pdf(self, knowledge_base_name: str, pdf_file_path: str, dpi: int = 200):
         print("[1/2] rendering pdf to images..")
         images = pdf_to_images(pdf_file_path, dpi=dpi)
         print("[2/2] model encoding images..")
@@ -94,10 +94,10 @@ if __name__ == "__main__":
     retriever = PDFVisualRetrieval(model=model, tokenizer=tokenizer)
     retriever.add_pdf('test', pdf_path)
     
-    topk_doc_ids_np, topk_values_np, images_topk = retriever.retrieve(knowledge_base='test', query='what is the number of VQ of this kind of codec method?', topk=1)
+    topk_doc_ids_np, topk_values_np, images_topk = retriever.retrieve(knowledge_base='test', query='what is the number of VQ of this kind of codec method?', topk=5)
     # 2
-    topk_doc_ids_np, topk_values_np, images_topk = retriever.retrieve(knowledge_base='test', query='the training loss curve of this paper?', topk=1)
+    topk_doc_ids_np, topk_values_np, images_topk = retriever.retrieve(knowledge_base='test', query='the training loss curve of this paper?', topk=5)
     # 3
-    topk_doc_ids_np, topk_values_np, images_topk = retriever.retrieve(knowledge_base='test', query='the experiment table?', topk=1)
+    topk_doc_ids_np, topk_values_np, images_topk = retriever.retrieve(knowledge_base='test', query='the experiment table?', topk=5)
     # 2