Update app.py

app.py

@@ -1,56 +1,134 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
+
 import tqdm
 from PIL import Image
+import hashlib
 import torch
 import fitz
+import threading
 import gradio as gr
 import spaces
 import os
 from transformers import AutoModel
 from transformers import AutoTokenizer
 import numpy as np
+import json
 
-cache_dir = 'pdf_cache'
+cache_dir = '/data/kb_cache'
 os.makedirs(cache_dir, exist_ok=True)
 
+def get_image_md5(img: Image.Image):
+    img_byte_array = img.tobytes()
+    hash_md5 = hashlib.md5()
+    hash_md5.update(img_byte_array)
+    hex_digest = hash_md5.hexdigest()
+    return hex_digest
+
+def calculate_md5_from_binary(binary_data):
+    hash_md5 = hashlib.md5()
+    hash_md5.update(binary_data)
+    return hash_md5.hexdigest()
+
+@spaces.GPU(duration=100)
+def add_pdf_gradio(pdf_file_binary, progress=gr.Progress()):
+    global model, tokenizer
+    model.eval()
+    
+    this_cache_dir = os.path.join(cache_dir, 'temp_cache')
+    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 = 200
+    doc = fitz.open("pdf", pdf_file_binary)
+    
+    reps_list = []
+    images = []
+    image_md5s = []
+
+    for page in progress.tqdm(doc):
+        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():
+            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')
+    
+    return "PDF processed successfully!"
+
+@spaces.GPU(duration=50)
+def retrieve_gradio(query: str, topk: int):
+    global model, tokenizer
+
+    model.eval()
+
+    target_cache_dir = os.path.join(cache_dir, 'temp_cache')
+
+    if not os.path.exists(target_cache_dir):
+        return None
+    
+    md5s = []
+    with open(os.path.join(target_cache_dir, f"md5s.txt"), 'r') as f:
+        for line in f:
+            md5s.append(line.rstrip('\n'))
+    
+    doc_reps = np.load(os.path.join(target_cache_dir, f"reps.npy"))
+
+    query_with_instruction = "Represent this query for retrieving relevant document: " + query
+    with torch.no_grad():
+        query_rep = model(text=[query_with_instruction], image=[None], tokenizer=tokenizer).reps.squeeze(0).cpu()
+
+    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_doc_ids_np = topk_doc_ids.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' 
 
-print("Embedding model loading...")
-model_path = 'RhapsodyAI/minicpm-visual-embedding-v0'
+print("emb model load begin...")
+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.eval()
 model.to(device)
-print("Embedding model loaded successfully!")
+print("emb model load success!")
 
-print("Generation model loading...")
+print("gen model load begin...")
 gen_model_path = 'openbmb/MiniCPM-V-2_6'
 gen_tokenizer = AutoTokenizer.from_pretrained(gen_model_path, trust_remote_code=True)
 gen_model = AutoModel.from_pretrained(gen_model_path, trust_remote_code=True, attn_implementation='sdpa', torch_dtype=torch.bfloat16)
 gen_model.eval()
 gen_model.to(device)
-print("Generation model loaded successfully!")
-
-@spaces.GPU(duration=100)
-def process_pdf(pdf_file, max_pages, progress=gr.Progress()):
-    doc = fitz.open("pdf", pdf_file)
-    num_pages = min(max_pages, len(doc))
-    
-    images = []
-    for page_num in progress.tqdm(range(num_pages)):
-        page = doc[page_num]
-        pix = page.get_pixmap(dpi=200)
-        image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-        images.append(image)
-    
-    return images
+print("gen model load success!")
 
 @spaces.GPU(duration=50)
 def answer_question(images, question):
     global gen_model, gen_tokenizer
-    images_ = [img.convert('RGB') for img in images]
+    images_ = [Image.open(image[0]).convert('RGB') for image in images]
     msgs = [{'role': 'user', 'content': [question, *images_]}]
     answer = gen_model.chat(
         image=None,
@@ -61,44 +139,43 @@ def answer_question(images, question):
     return answer
 
 with gr.Blocks() as app:
-    gr.Markdown("# PDF Question Answering with Vision Language Model")
-
-    gr.Markdown("""
-    This application uses a Vision Language Model to answer questions about PDF documents.
+    gr.Markdown("# MiniCPMV-RAG-PDFQA: Two Vision Language Models Enable End-to-End RAG")
     
-    1. Upload a PDF file
-    2. Set the maximum number of pages to process
-    3. Click "Process PDF" to extract the pages
-    4. Enter your question about the PDF content
-    5. Click "Answer Question" to get the model's response
-    """)
+    gr.Markdown("""
+- A Vision Language Model Dense Retriever ([minicpm-visual-embedding-v0](https://huggingface.co/RhapsodyAI/minicpm-visual-embedding-v0)) **directly reads** your PDFs **without need of OCR**, produce **multimodal dense representations** and build your personal library. 
+
+- **Ask a question**, it retrieves the most relevant pages, then [MiniCPM-V-2.6](https://huggingface.co/spaces/openbmb/MiniCPM-V-2_6) will answer your question based on pages recalled, with strong multi-image understanding capability.
 
+    - It helps you read a long **visually-intensive** or **text-oriented** PDF document and find the pages that answer your question.
+
+    - It helps you build a personal library and retrieve book pages from a large collection of books.
+
+    - It works like a human: read, store, retrieve, and answer with full vision.
+""")
+    
+    gr.Markdown("- Currently online demo support PDF document with less than 50 pages due to GPU time limit. Deploy on your own machine for longer PDFs and books.")
+    
     with gr.Row():
-        file_input = gr.File(type="binary", label="Upload PDF")
-        max_pages = gr.Number(value=10, minimum=1, maximum=50, step=1, label="Maximum number of pages to process")
+        file_input = gr.File(type="binary", label="Step 1: Upload PDF")
         process_button = gr.Button("Process PDF")
+        file_result = gr.Textbox(label="PDF Process Status")
+    
+    process_button.click(add_pdf_gradio, inputs=[file_input], outputs=file_result)
 
     with gr.Row():
         query_input = gr.Text(label="Your Question")
-        answer_button = gr.Button("Answer Question")
-
-    images_output = gr.Gallery(label="Processed PDF Pages", visible=False)
-    gen_model_response = gr.Textbox(label="Model's Answer")
-
-    def process_and_show(pdf_file, max_pages):
-        images = process_pdf(pdf_file, max_pages)
-        return gr.Gallery.update(value=images, visible=True)
+        topk_input = gr.Number(value=5, minimum=1, maximum=10, step=1, label="Number of Pages to Retrieve")
+        retrieve_button = gr.Button("Retrieve Pages")
+        images_output = gr.Gallery(label="Retrieved Pages")
+    
+    retrieve_button.click(retrieve_gradio, inputs=[query_input, topk_input], outputs=images_output)
 
-    process_button.click(
-        process_and_show,
-        inputs=[file_input, max_pages],
-        outputs=images_output
-    )
+    with gr.Row():
+        answer_button = gr.Button("Answer Question")
+        gen_model_response = gr.Textbox(label="MiniCPM-V-2.6's Answer")
 
-    answer_button.click(
-        answer_question,
-        inputs=[images_output, query_input],
-        outputs=gen_model_response
-    )
+        answer_button.click(fn=answer_question, inputs=[images_output, query_input], outputs=gen_model_response)
+    
+    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()
+app.launch()