app.py

#!/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 = '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("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("emb model load success!")

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("gen model load success!")

@spaces.GPU(duration=50)
def answer_question(images, question):
    global gen_model, gen_tokenizer
    images_ = [Image.open(image[0]).convert('RGB') for image in images]
    msgs = [{'role': 'user', 'content': [question, *images_]}]
    answer = gen_model.chat(
        image=None,
        msgs=msgs,
        tokenizer=gen_tokenizer
    )
    print(answer)
    return answer

with gr.Blocks() as app:
    gr.Markdown("# MiniCPMV-RAG-PDFQA: Two Vision Language Models Enable End-to-End RAG")
    
    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="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")
        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)

    with gr.Row():
        answer_button = gr.Button("Answer Question")
        gen_model_response = gr.Textbox(label="MiniCPM-V-2.6's Answer")

        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()