"""Template Demo for IBM Granite Hugging Face spaces.""" import os import time from pathlib import Path import re import gradio as gr import spaces import torch from gradio_pdf import PDF from sandbox.light_rag.light_rag import LightRAG from themes.research_monochrome import theme dir_ = Path(__file__).parent.parent TITLE = "Multimodal RAG with Granite Vision 3.2" DESCRIPTION = """

This experimental demo highlights granite-vision-3.2-2b capabilities within a multimodal retrieval-augmented generation (RAG) pipeline, demonstrating Granite's document understanding in real-world applications. Explore the sample document excerpts and try the sample prompts or enter your own. Keep in mind that AI can occasionally make mistakes. View Documentation

""" device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu") BASE_PATH = dir_ / "data" / "final_v2_mar04" PDFS_PATH = BASE_PATH / "pdfs" MILVUS_PATH = BASE_PATH / "milvus" IMAGES_PATH = BASE_PATH / "images" PREVIEWS_PATH = BASE_PATH / "preview" sample_data = [ { "preview_image": str(PREVIEWS_PATH / "IBM-financial-2010.png"), "prompts": """Where geographically was the greatest growth in revenue in 2007? Which year had the highest income in billion? Did the net income decrease in 2007 compared to 2006? Net cash from operations on 2005? What does it mean to be Globally Integrated Enterprise? What are the segments for pretax income?""".split("\n"), "pdf": str(PDFS_PATH / "IBM_Annual_Report_2007_3-20.pdf"), "index": "ibm_report_2007_short_text_milvus_lite_2048_128_slate_278m_cosine", "db": str(MILVUS_PATH / "milvus.db"), "name": "IBM annual report 2007", "origin": "https://www.ibm.com/investor/att/pdf/IBM_Annual_Report_2007.pdf", "image_paths": {"prefix": str(IMAGES_PATH / "ibm_report_2007") + "/", "use_last": 2}, }, { "preview_image": str(PREVIEWS_PATH / "Wilhlborg-financial.png"), "prompts": """Where does Wihlborgs mainly operate? Which year had the second lowest Equity/assets ratio? Which year had the highest Project investments value? What is the trend of equity/assets ratio? What was the Growth percentage in income from property management in 2020? Has the company’s interest coverage ratio increased or decreased in recent years?""".split("\n") , "pdf": str(PDFS_PATH / "wihlborgs-2-13_16-18.pdf"), "index": "wihlborgs_short_text_milvus_lite_2048_128_slate_278m_cosine", "db": str(MILVUS_PATH / "milvus.db"), "name": "Wihlborgs Report 2020", "origin": "https://www.wihlborgs.se/globalassets/investor-relations/rapporter/2021/20210401-wihlborgs-annual-report-and-sustainability-report-2020-c24a6b51-c124-44fc-a4af-4237a33a29fb.pdf", "image_paths": {"prefix": str(IMAGES_PATH / "wihlborgs") + "/", "use_last": 2}, }, ] config = { "embedding_model_id": "ibm-granite/granite-embedding-278m-multilingual", "generation_model_id": "ibm-granite/granite-3.1-8b-instruct", "milvus_collection_name": "granite_vision_tech_report_text_milvus_lite_512_128_slate_125m_cosine", "milvus_db_path": str(dir_ / "data" / MILVUS_PATH / "milvus_text_sample.db"), } if gr.NO_RELOAD: light_rag: LightRAG = LightRAG(config) if not os.environ.get("LAZY_LOADING") == "true": for sample in sample_data: light_rag.precache_milvus(sample["index"], sample["db"]) def lower_md_headers(md: str) -> str: return re.sub(r'(?:^|\n)##?\s(.+)', lambda m: '\n### ' + m.group(1), md) # Parser for retrival results def format_retrieval_result(i, d, cb, selected_sample): image_paths = sample_data[selected_sample]["image_paths"] if d.metadata["type"] == "text": context_string = f"---\n## Context {i + 1}\n#### (text extracted from document)\n{lower_md_headers(d.page_content)}\n" cb.append(gr.ChatMessage(role="assistant", content=context_string)) return True elif d.metadata["type"] == "image_description": context_string = f"---\n## Context {i + 1}\n#### (image description generated by Granite Vision)" cb.append(gr.ChatMessage(role="assistant", content=context_string)) # /dccstor/mm-rag/idanfr/granite_vision_demo/wdu_output/IBM_Annual_Report_2007/images/IBM_Annual_Report_2007_im_image_7_1.png image_path_parts = d.metadata["image_fullpath"].split("/") image_path = image_paths["prefix"] + ("/".join(image_path_parts[-image_paths["use_last"]:])) # print(f"image_path: {image_path}") cb.append(gr.ChatMessage(role="assistant", content=gr.Image(image_path))) cb.append(gr.ChatMessage(role="assistant", content=f"\n{lower_md_headers(d.metadata['image_description'])}\n")) chatbot = gr.Chatbot( examples=[{"text": x} for x in sample_data[0]["prompts"]], type="messages", label=f"Q&A about {sample_data[0]['name']}", height=685, group_consecutive_messages=True, autoscroll=False, elem_classes=["chatbot_view"], ) @spaces.GPU() def generate_with_llm(query, context): if os.environ.get("NO_LLM"): time.sleep(2) return "Now answer, just a string", query return light_rag.generate(query=query, context=context) # TODO: maybe add GPU back ? def retrieval(collection, db, q): return light_rag.search(q, top_n=3, collection=collection, db=db) # ################ # User Interface # ################ css_file_path = Path(Path(__file__).parent / "app.css") head_file_path = Path(Path(__file__).parent / "app_head.html") with gr.Blocks(fill_height=True, css_paths=css_file_path, head_paths=head_file_path, theme=theme, title=TITLE) as demo: is_in_edit_mode = gr.State(True) # in block to be reactive selected_doc = gr.State(0) current_question = gr.State("") gr.Markdown(f"# {TITLE}") gr.Markdown(DESCRIPTION) with gr.Row(): # LEFT COLUMN: Sample selection, download, and PDF viewer with gr.Column(): # Show preview images images_only = [sd["preview_image"] for sd in sample_data] document_gallery = gr.Gallery( images_only, label="Select a document", rows=1, columns=3, height="125px", # width="125px", allow_preview=False, selected_index=0, elem_classes=["preview_im_element"], ) with gr.Group(): pdf_display = PDF( sample_data[0]["pdf"], label=f"Preview for {sample_data[0]['name']}", height=460, interactive=False, elem_classes=["pdf_viewer"], ) dl_btn = gr.DownloadButton( label=f"Download PDF ({sample_data[0]['name']})", value=sample_data[0]["pdf"], visible=True ) def sample_image_selected(d: gr.SelectData): dx = sample_data[d.index] # print(f"DX:{dx}") return ( gr.update(examples=[{"text": x} for x in dx["prompts"]], label=f"Q&A about {dx['name']}"), gr.update(value=dx["pdf"], label=f"Preview for {dx['name']}"), gr.DownloadButton(value=dx["pdf"], label=f"Download PDF ({dx['name']})"), d.index ) document_gallery.select(lambda: [], outputs=[chatbot]) document_gallery.select(sample_image_selected, inputs=[], outputs=[chatbot, pdf_display, dl_btn, selected_doc]) # Right Column: Chat interface with gr.Column(): # Render ChatBot chatbot.render() # Define behavior for example selection def update_user_chat_x(x: gr.SelectData): return [gr.ChatMessage(role="user", content=x.value["text"])] def question_from_selection(x: gr.SelectData): return x.value["text"] def _decorate_yield_result(cb, fb_status=False, gallery_status=False): return cb, gr.Button(interactive=fb_status), gr.Gallery( elem_classes=["preview_im_element"] if gallery_status else ["preview_im_element", "inactive_div"]) def send_generate(msg, cb, selected_sample): collection = sample_data[selected_sample]["index"] db = sample_data[selected_sample]["db"] # print(f"collection: {collection}, {db}") original_msg = gr.ChatMessage(role="user", content=msg) cb.append(original_msg) waiting_for_retrieval_msg = gr.ChatMessage(role="assistant", content='## Answer\n*Querying Index*. . .') cb.append(waiting_for_retrieval_msg) yield _decorate_yield_result(cb) q = msg.strip() results = retrieval(collection, db, q) # for d in results: # print(f"****\n{d}") context_string = "## Context Documents for Answer\n\n" for i, d in enumerate(results): if format_retrieval_result(i, d, cb, selected_sample): yield _decorate_yield_result(cb) waiting_for_llm_msg = gr.ChatMessage(role="assistant", content='## Answer\n *Waiting for LLM* . . . ') cb[1] = waiting_for_llm_msg yield _decorate_yield_result(cb) answer, prompt = generate_with_llm(q, results) cb[1] = gr.ChatMessage(role="assistant", content=f"## Answer\n{answer.strip()}") # cb.pop() # cb.append(gr.ChatMessage(role="assistant", content=f"## Answer\n{answer.strip()}")) yield _decorate_yield_result(cb, fb_status=True, gallery_status=True) # Create User Chat Textbox and Reset Button tbb = gr.Textbox(submit_btn=True, show_label=False, placeholder="Type a message...") fb = gr.Button("Ask new question", visible=False) fb.click(lambda: [], outputs=[chatbot]) chatbot.example_select(lambda: False, outputs=is_in_edit_mode) # chatbot.example_select(update_user_chat_x, outputs=[chatbot]) chatbot.example_select(question_from_selection, inputs=[], outputs=[current_question] ).then(send_generate, inputs=[current_question, chatbot, selected_doc], outputs=[chatbot, fb, document_gallery]) def textbox_switch(e_mode): # Handle toggling between edit and non-edit mode if not e_mode: return [gr.update(visible=False), gr.update(visible=True)] else: return [gr.update(visible=True), gr.update(visible=False)] tbb.submit(lambda: False, outputs=[is_in_edit_mode]) fb.click(lambda: True, outputs=[is_in_edit_mode]) is_in_edit_mode.change(textbox_switch, inputs=[is_in_edit_mode], outputs=[tbb, fb]) # submit user question # tbb.submit(lambda x: [gr.ChatMessage(role="user", content=x)], inputs=tbb, outputs=chatbot) tbb.submit(lambda x: x, inputs=[tbb], outputs=[current_question] ).then(send_generate, inputs=[current_question, chatbot, selected_doc], outputs=[chatbot, fb, document_gallery]) if __name__ == "__main__": # demo.queue(max_size=20).launch() demo.launch()