finish

app.py

@@ -16,7 +16,17 @@ import fitz
 import threading
 import gradio as gr
 import spaces
+import os
+from transformers import AutoModel
+from transformers import AutoTokenizer
+from PIL import Image
+import torch
+import os
+import numpy as np
+import json
 
+cache_dir = '/data/kb_cache'
+os.makedirs(cache_dir, exist_ok=True)
 
 def get_image_md5(img: Image.Image):
     img_byte_array = img.tobytes()
@@ -25,152 +35,129 @@ def get_image_md5(img: Image.Image):
     hex_digest = hash_md5.hexdigest()
     return hex_digest
 
-def pdf_to_images(pdf_path, dpi=100):
-    doc = fitz.open(pdf_path)
-    images = []
-    for page in tqdm.tqdm(doc):
-        pix = page.get_pixmap(dpi=dpi)
-        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-        images.append(img)
-    return images
-
 def calculate_md5_from_binary(binary_data):
     hash_md5 = hashlib.md5()
     hash_md5.update(binary_data)
     return hash_md5.hexdigest()
 
-class PDFVisualRetrieval:
-    def __init__(self, model, tokenizer):
-        self.tokenizer = tokenizer
-        self.model = model
-        self.reps = {}
-        self.images = {}
-        self.lock = threading.Lock()
+@spaces.GPU(duration=120)
+def add_pdf_gradio(pdf_file_binary, progress=gr.Progress()):
+    global model, tokenizer
+
+    knowledge_base_name = calculate_md5_from_binary(pdf_file_binary)
     
-    def retrieve(self, knowledge_base: str, query: str, topk: int):
-        doc_reps = list(self.reps[knowledge_base].values())
-        query_with_instruction = "Represent this query for retrieving relavant document: " + query
-        with torch.no_grad():
-            query_rep = self.model(text=[query_with_instruction], image=[None], tokenizer=self.tokenizer).reps.squeeze(0)
-        doc_reps_cat = torch.stack(doc_reps, dim=0)
-        similarities = torch.matmul(query_rep, doc_reps_cat.T)
-        topk_values, topk_doc_ids = torch.topk(similarities, k=topk)
-        topk_values_np = topk_values.cpu().numpy()
-        topk_doc_ids_np = topk_doc_ids.cpu().numpy()
-        similarities_np = similarities.cpu().numpy()
-        all_images_doc_list = list(self.images[knowledge_base].values())
-        images_topk = [all_images_doc_list[idx] for idx in topk_doc_ids_np]
-        return topk_doc_ids_np, topk_values_np, images_topk
+    this_cache_dir = os.path.join(cache_dir, knowledge_base_name)
+    os.makedirs(this_cache_dir, exist_ok=True)
+
+    with open(os.path.join(this_cache_dir, f"src.pdf"), 'wb') as file:
+        file.write(pdf_file_binary)
+
+    dpi = 100
+    doc = fitz.open("pdf", pdf_file_binary)
     
-    def add_pdf(self, knowledge_base_name: str, pdf_file_path: str, dpi: int = 100):
-        if knowledge_base_name not in self.reps:
-            self.reps[knowledge_base_name] = {}
-        if knowledge_base_name not in self.images:
-            self.images[knowledge_base_name] = {}
-        doc = fitz.open(pdf_file_path)
-        print("model encoding images..")
-        for page in tqdm.tqdm(doc):
-            pix = page.get_pixmap(dpi=dpi)
-            image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-            image_md5 = get_image_md5(image)
-            with torch.no_grad():
-                reps = self.model(text=[''], image=[image], tokenizer=self.tokenizer).reps
-            self.reps[knowledge_base_name][image_md5] = reps.squeeze(0)
-            self.images[knowledge_base_name][image_md5] = image
-        return
-
-    def add_pdf_gradio(self, pdf_file_binary, progress=gr.Progress()):
-        knowledge_base_name = calculate_md5_from_binary(pdf_file_binary)
-        if knowledge_base_name not in self.reps:
-            self.reps[knowledge_base_name] = {}
-        else:
-            return knowledge_base_name
-        if knowledge_base_name not in self.images:
-            self.images[knowledge_base_name] = {}
-        dpi = 100
-        doc = fitz.open("pdf", pdf_file_binary)
-        
-        for page in progress.tqdm(doc):
-            # with self.lock: # because we hope one 16G gpu only process one image at the same time
-            pix = page.get_pixmap(dpi=dpi)
-            image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-            image_md5 = get_image_md5(image)
-            with torch.no_grad():
-                reps = self.model(text=[''], image=[image], tokenizer=self.tokenizer).reps
-            self.reps[knowledge_base_name][image_md5] = reps.squeeze(0)
-            self.images[knowledge_base_name][image_md5] = image
-        
-        return knowledge_base_name
-
-    def retrieve_gradio(self, knowledge_base: str, query: str, topk: int):
-        doc_reps = list(self.reps[knowledge_base].values())
-        query_with_instruction = "Represent this query for retrieving relavant document: " + query
+    reps_list = []
+    images = []
+    image_md5s = []
+
+    for page in progress.tqdm(doc):
+        # with self.lock: # because we hope one 16G gpu only process one image at the same time
+        pix = page.get_pixmap(dpi=dpi)
+        image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+        image_md5 = get_image_md5(image)
+        image_md5s.append(image_md5)
         with torch.no_grad():
-            query_rep = self.model(text=[query_with_instruction], image=[None], tokenizer=self.tokenizer).reps.squeeze(0)
-        doc_reps_cat = torch.stack(doc_reps, dim=0)
-        similarities = torch.matmul(query_rep, doc_reps_cat.T)
-        topk_values, topk_doc_ids = torch.topk(similarities, k=topk)
-        topk_values_np = topk_values.cpu().numpy()
-        topk_doc_ids_np = topk_doc_ids.cpu().numpy()
-        similarities_np = similarities.cpu().numpy()
-        all_images_doc_list = list(self.images[knowledge_base].values())
-        images_topk = [all_images_doc_list[idx] for idx in topk_doc_ids_np]
-        return images_topk
-
-
-if __name__ == "__main__":
-    from transformers import AutoModel
-    from transformers import AutoTokenizer
-    from PIL import Image
-    import torch
+            reps = model(text=[''], image=[image], tokenizer=tokenizer).reps
+        reps_list.append(reps.squeeze(0).cpu().numpy())
+        images.append(image)
+
+    for idx in range(len(images)):
+        image = images[idx]
+        image_md5 = image_md5s[idx]
+        cache_image_path = os.path.join(this_cache_dir, f"{image_md5}.png")
+        image.save(cache_image_path)
+
+    np.save(os.path.join(this_cache_dir, f"reps.npy"), reps_list)
+
+    with open(os.path.join(this_cache_dir, f"md5s.txt"), 'w') as f:
+        for item in image_md5s:
+            f.write(item+'\n')
     
-    device = 'cuda:0'
+    return knowledge_base_name
+
+# @spaces.GPU
+def retrieve_gradio(knowledge_base: str, query: str, topk: int):
+    global model, tokenizer
+
+    target_cache_dir = os.path.join(cache_dir, knowledge_base)
+
+    if not os.path.exists(target_cache_dir):
+        return None
     
-    # Load model, be sure to substitute `model_path` by your model path 
-    model_path = 'RhapsodyAI/minicpm-visual-embedding-v0' # replace with your local model path
-    # pdf_path = "/home/jeeves/xubokai/minicpm-visual-embedding-v0/2406.07422v1.pdf"
+    md5s = []
+    with open(os.path.join(target_cache_dir, f"md5s.txt"), 'r') as f:
+        for line in f:
+            md5s.append(line.rstrip('\n'))
     
-    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-    model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
-    model.to(device)
+    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():
+        query_rep = model(text=[query_with_instruction], image=[None], tokenizer=tokenizer).reps.squeeze(0).cpu()
+
+    query_md5 = hashlib.md5(query.encode()).hexdigest()
+    with open(os.path.join(target_cache_dir, f"q-{query_md5}.json"), 'w') as f:
+        f.write(json.dumps(
+            {
+                "query": query
+            }, indent=4, ensure_ascii=False
+        ))
+
+    doc_reps_cat = torch.stack([torch.Tensor(i) for i in doc_reps], dim=0)
+
+    similarities = torch.matmul(query_rep, doc_reps_cat.T)
+
+    topk_values, topk_doc_ids = torch.topk(similarities, k=topk)
+
+    topk_values_np = topk_values.cpu().numpy()
+
+    topk_doc_ids_np = topk_doc_ids.cpu().numpy()
+
+    similarities_np = similarities.cpu().numpy()
+
+    images_topk = [Image.open(os.path.join(target_cache_dir, f"{md5s[idx]}.png")) for idx in topk_doc_ids_np]
+
+    return images_topk
+
+
+device = 'cuda' 
+model_path = 'RhapsodyAI/minicpm-visual-embedding-v0' # replace with your local model path
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
+model.to(device)
+
+
+with gr.Blocks() as app:
+    gr.Markdown("# Memex: OCR-free Visual Document Retrieval @RhapsodyAI")
     
-    retriever = PDFVisualRetrieval(model=model, tokenizer=tokenizer)
+    with gr.Row():
+        file_input = gr.File(type="binary", label="Upload PDF")
+        file_result = gr.Text(label="Knowledge Base ID (remember this!)")
+        process_button = gr.Button("Process PDF")
     
-    @spaces.GPU
-    def add_pdf_gradio(pdf_file_binary):
-        return retriever.add_pdf_gradio(pdf_file_binary)
+    process_button.click(add_pdf_gradio, inputs=[file_input], outputs=file_result)
+
+    with gr.Row():
+        kb_id_input = gr.Text(label="Your Knowledge Base ID")
+        query_input = gr.Text(label="Your Queston")
+        topk_input = inputs=gr.Number(value=1, minimum=1, maximum=5, step=1, label="Top K")
+        retrieve_button = gr.Button("Retrieve")
     
-    @spaces.GPU
-    def retrieve_gradio(knowledge_base, query, topk):
-        return retriever.retrieve_gradio(knowledge_base, query, topk)
-
-    # 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)
-    # # 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)
-    # # 3
-    # topk_doc_ids_np, topk_values_np, images_topk = retriever.retrieve(knowledge_base='test', query='the experiment table?', topk=1)
-    # # 2
+    with gr.Row():
+        images_output = gr.Gallery(label="Retrieved Pages")
     
-    with gr.Blocks() as app:
-        gr.Markdown("# Memex: OCR-free Visual Document Retrieval @RhapsodyAI")
-        
-        with gr.Row():
-            file_input = gr.File(type="binary", label="Upload PDF")
-            file_result = gr.Text(label="Knowledge Base ID (remember this!)")
-            process_button = gr.Button("Process PDF")
-        
-        process_button.click(add_pdf_gradio, inputs=[file_input], outputs=file_result)
-
-        with gr.Row():
-            kb_id_input = gr.Text(label="Your Knowledge Base ID")
-            query_input = gr.Text(label="Your Queston")
-            topk_input = inputs=gr.Number(value=1, minimum=1, maximum=5, step=1, label="Top K")
-            retrieve_button = gr.Button("Retrieve")
-        
-        with gr.Row():
-            images_output = gr.Gallery(label="Retrieved Pages")
-        
-        retrieve_button.click(retrieve_gradio, inputs=[kb_id_input, query_input, topk_input], outputs=images_output)
-
-    app.launch()
+    retrieve_button.click(retrieve_gradio, inputs=[kb_id_input, query_input, topk_input], outputs=images_output)
+
+    gr.Markdown("By using this demo, you agree to share your use data with us for research purpose, to help improve user experience.")
+
+app.launch()