Update app.py

app.py

@@ -15,7 +15,7 @@ import os
 import numpy as np
 import json
 
-cache_dir = '/home/user/data'
+cache_dir = '/data/KB'
 os.makedirs(cache_dir, exist_ok=True)
 
 def get_image_md5(img: Image.Image):
@@ -33,7 +33,8 @@ def calculate_md5_from_binary(binary_data):
 @spaces.GPU(duration=100)
 def add_pdf_gradio(pdf_file_binary, progress=gr.Progress()):
     global model, tokenizer
-
+    model.eval()
+    
     knowledge_base_name = calculate_md5_from_binary(pdf_file_binary)
     
     this_cache_dir = os.path.join(cache_dir, knowledge_base_name)
@@ -78,6 +79,8 @@ def add_pdf_gradio(pdf_file_binary, progress=gr.Progress()):
 def retrieve_gradio(knowledge_base: str, query: str, topk: int):
     global model, tokenizer
 
+    model.eval()
+
     target_cache_dir = os.path.join(cache_dir, knowledge_base)
 
     if not os.path.exists(target_cache_dir):
@@ -90,7 +93,7 @@ def retrieve_gradio(knowledge_base: str, query: str, topk: int):
     
     doc_reps = np.load(os.path.join(target_cache_dir, f"reps.npy"))
 
-    query_with_instruction = "Represent this query for retrieving relevant document: " + query
+    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()
 
@@ -166,53 +169,91 @@ def downvote(knowledge_base, query):
     return
 
 
+
 device = 'cuda' 
-model_path = 'RhapsodyAI/minicpm-visual-embedding-v0' # replace with your local model path
+
+print("emb model load begin...")
+model_path = 'openbmb/VisRAG-Ret' # 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, attn_implementation='sdpa', trust_remote_code=True)
+gen_model =  AutoModel.from_pretrained(gen_model_path, trust_remote_code=True, 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
+    # here each element of images is a tuple of (image_path, None).
+    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("# Memex: OCR-free Visual Document Embedding Model as Your Personal Librarian")
-    gr.Markdown("""The model only takes images as document-side inputs and produce vectors representing document pages. Memex is trained with over 200k query-visual document pairs, including textual document, visual document, arxiv figures, plots, charts, industry documents, textbooks, ebooks, and openly-available PDFs, etc. Its performance is on a par with our ablation text embedding model on text-oriented documents, and an advantages on visually-intensive documents.
-Our model is capable of:
-- Help you read a long visually-intensive or text-oriented PDF document and find the pages that answer your question.
-- Help you build a personal library and retireve book pages from a large collection of books.
-- It works like human: read and comprehend with vision and remember multimodal information in hippocampus.""")
-
-    gr.Markdown("- Our model is proudly based on MiniCPM-V series [MiniCPM-V-2.6](https://huggingface.co/openbmb/MiniCPM-V-2_6) [MiniCPM-V-2](https://huggingface.co/openbmb/MiniCPM-V-2).")
+    gr.Markdown("# MiniCPMV-RAG-PDFQA: Two Vision Language Models Enable End-to-End RAG")
     
-    gr.Markdown("- We open-sourced our model at [RhapsodyAI/minicpm-visual-embedding-v0](https://huggingface.co/RhapsodyAI/minicpm-visual-embedding-v0)")
+    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 retrieve most relavant 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 retireve book pages from a large collection of books.
+
+    - It works like a human: read, store, retrieve, and answer with full vision.
+""")
     
-    gr.Markdown("- Currently we support PDF document with less than 50 pages, PDF over 50 pages will reach GPU time limit.")
+    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")
-        file_result = gr.Text(label="Knowledge Base ID (remember this!)")
+        file_input = gr.File(type="binary", label="Step 1: Upload PDF")
+        file_result = gr.Text(label="Knowledge Base ID (remember it, it is re-usable!)")
         process_button = gr.Button("Process PDF (Don't click until PDF upload success)")
     
     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 (paste your Knowledge Base ID here:)")
+        kb_id_input = gr.Text(label="Your Knowledge Base ID (paste your Knowledge Base ID here, it is re-usable:)")
         query_input = gr.Text(label="Your Queston")
         topk_input = inputs=gr.Number(value=5, minimum=1, maximum=10, step=1, label="Number of pages to retrieve")
-        retrieve_button = gr.Button("Step 1: Retrieve")
-    
-    with gr.Row():
-        downvote_button = gr.Button("🤣Downvote")
-        upvote_button = gr.Button("🤗Upvote")
+        retrieve_button = gr.Button("Step2: Retrieve Pages")
     
     with gr.Row():
-        images_output = gr.Gallery(label="Step 2: Retrieved Pages")
+        images_output = gr.Gallery(label="Retrieved Pages")
     
     retrieve_button.click(retrieve_gradio, inputs=[kb_id_input, query_input, topk_input], outputs=images_output)
 
+    with gr.Row():
+        button = gr.Button("Step 3: Answer Question with Retrieved Pages")
+
+        gen_model_response = gr.Textbox(label="MiniCPM-V-2.6's Answer")
+
+        button.click(fn=answer_question, inputs=[images_output, query_input], outputs=gen_model_response)
+    
+    with gr.Row():
+        downvote_button = gr.Button("🤣Downvote")
+        upvote_button = gr.Button("🤗Upvote")
+        
     upvote_button.click(upvote, inputs=[kb_id_input, query_input], outputs=None)
     downvote_button.click(downvote, inputs=[kb_id_input, query_input], outputs=None)
 
     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()