Create app.py

app.py

@@ -0,0 +1,227 @@
+import gradio as gr
+import os
+import nest_asyncio
+import re
+from pathlib import Path
+import typing as t
+import base64
+from mimetypes import guess_type
+from llama_parse import LlamaParse
+from llama_index.core.schema import TextNode
+from llama_index.core import VectorStoreIndex, StorageContext, load_index_from_storage, Settings
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.llms.openai import OpenAI
+from llama_index.core.query_engine import CustomQueryEngine
+from llama_index.multi_modal_llms.openai import OpenAIMultiModal
+from llama_index.core.prompts import PromptTemplate
+from llama_index.core.schema import ImageNode
+from llama_index.core.base.response.schema import Response
+from typing import Any, List, Optional
+from llama_index.core.postprocessor.types import BaseNodePostprocessor
+
+nest_asyncio.apply()
+
+# Setting API keys
+os.environ["OPENAI_API_KEY"] = os.getenv('OPENAI_API_KEY')
+os.environ["LLAMA_CLOUD_API_KEY"] = os.getenv('LLAMA_CLOUD_API_KEY')
+
+# Initialize the parser
+parser = LlamaParse(
+    result_type="markdown",
+    parsing_instruction="You are given a medical textbook on medicine",
+    use_vendor_multimodal_model=True,
+    vendor_multimodal_model_name="gpt-4o-mini-2024-07-18",
+    show_progress=True,
+    verbose=True,
+    invalidate_cache=True,
+    do_not_cache=True,
+    num_workers=8,
+    language="en"
+)
+
+# Function to encode image to data URL
+def local_image_to_data_url(image_path):
+    mime_type, _ = guess_type(image_path)
+    if mime_type is None:
+        mime_type = 'image/png'
+    with open(image_path, "rb") as image_file:
+        base64_encoded_data = base64.b64encode(image_file.read()).decode('utf-8')
+    return f"data:{mime_type};base64,{base64_encoded_data}"
+
+# Function to get sorted image files
+def get_page_number(file_name):
+    match = re.search(r"-page-(\d+)\.jpg$", str(file_name))
+    if match:
+        return int(match.group(1))
+    return 0
+
+def _get_sorted_image_files(image_dir):
+    raw_files = [f for f in list(Path(image_dir).iterdir()) if f.is_file()]
+    sorted_files = sorted(raw_files, key=get_page_number)
+    return sorted_files
+
+def get_text_nodes(md_json_objs, image_dir) -> t.List[TextNode]:
+    nodes = []
+    for result in md_json_objs:
+        json_dicts = result["pages"]
+        document_name = result["file_path"].split('/')[-1]
+        docs = [doc["md"] for doc in json_dicts]
+        image_files = _get_sorted_image_files(image_dir)
+        for idx, doc in enumerate(docs):
+            node = TextNode(
+                text=doc,
+                metadata={"image_path": str(image_files[idx]), "page_num": idx + 1, "document_name": document_name},
+            )
+            nodes.append(node)
+    return nodes
+
+# Gradio interface functions
+def upload_and_process_file(uploaded_file):
+    if uploaded_file is None:
+        return "Please upload a medical textbook (pdf)"
+    
+    file_path = f"{uploaded_file.name}"
+    with open(file_path, "wb") as f:
+        f.write(uploaded_file.read())
+
+    md_json_objs = parser.get_json_result([file_path])
+    image_dicts = parser.get_images(md_json_objs, download_path="data_images")
+    
+    return md_json_objs
+
+def ask_question(md_json_objs, query_text, uploaded_query_image=None):
+    if not md_json_objs:
+        return "No knowledge base loaded. Please upload a file first."
+    
+    text_nodes = get_text_nodes(md_json_objs, "data_images")
+
+    # Setup index and LLM
+    embed_model = OpenAIEmbedding(model="text-embedding-3-large")
+    llm = OpenAI("gpt-4o-mini-2024-07-18")
+    Settings.llm = llm
+    Settings.embed_model = embed_model
+
+    if not os.path.exists("storage_manuals"):
+        index = VectorStoreIndex(text_nodes, embed_model=embed_model)
+        index.storage_context.persist(persist_dir="./storage_manuals")
+    else:
+        ctx = StorageContext.from_defaults(persist_dir="./storage_manuals")
+        index = load_index_from_storage(ctx)
+
+    retriever = index.as_retriever()
+
+    # Encode query image if provided
+    encoded_image_url = None
+    if uploaded_query_image is not None:
+        query_image_path = f"{uploaded_query_image.name}"
+        with open(query_image_path, "wb") as img_file:
+            img_file.write(uploaded_query_image.read())
+        encoded_image_url = local_image_to_data_url(query_image_path)
+
+    # Setup query engine
+    QA_PROMPT_TMPL = """
+    You are a friendly medical chatbot designed to assist users by providing accurate and detailed responses to medical questions based on information from medical books.
+
+    ### Context:
+    ---------------------
+    {context_str}
+    ---------------------
+
+    ### Query Text:
+    {query_str}
+
+    ### Query Image:
+    ---------------------
+    {encoded_image_url}
+    ---------------------
+
+    ### Answer:
+    """
+    QA_PROMPT = PromptTemplate(QA_PROMPT_TMPL)
+    gpt_4o_mm = OpenAIMultiModal(model="gpt-4o-mini-2024-07-18")
+
+    class MultimodalQueryEngine(CustomQueryEngine):
+        qa_prompt: PromptTemplate
+        retriever: BaseRetriever
+        multi_modal_llm: OpenAIMultiModal
+        node_postprocessors: Optional[List[BaseNodePostprocessor]]
+
+        def __init__(
+            self,
+            qa_prompt: PromptTemplate,
+            retriever: BaseRetriever,
+            multi_modal_llm: OpenAIMultiModal,
+            node_postprocessors: Optional[List[BaseNodePostprocessor]] = [],
+        ):
+            super().__init__(
+                qa_prompt=qa_prompt,
+                retriever=retriever,
+                multi_modal_llm=multi_modal_llm,
+                node_postprocessors=node_postprocessors
+            )
+
+        def custom_query(self, query_str: str):
+            # retrieve most relevant nodes
+            nodes = self.retriever.retrieve(query_str)
+
+            # create image nodes from the image associated with those nodes
+            image_nodes = [
+                NodeWithScore(node=ImageNode(image_path=n.node.metadata["image_path"]))
+                for n in nodes
+            ]
+
+            # create context string from parsed markdown text
+            ctx_str = "\n\n".join(
+                [r.node.get_content(metadata_mode=MetadataMode.LLM).strip() for r in nodes]
+            )
+
+            # prompt for the LLM
+            fmt_prompt = self.qa_prompt.format(
+                context_str=ctx_str, query_str=query_str, encoded_image_url=encoded_image_url
+            )
+
+            # use the multimodal LLM to interpret images and generate a response to the prompt
+            llm_response = self.multi_modal_llm.complete(
+                prompt=fmt_prompt,
+                image_documents=[image_node.node for image_node in image_nodes],
+            )
+
+            return Response(
+                response=str(llm_response),
+                source_nodes=nodes,
+                metadata={"text_nodes": nodes, "image_nodes": image_nodes},
+            )
+
+    query_engine = MultimodalQueryEngine(QA_PROMPT, retriever, gpt_4o_mm)
+
+    response = query_engine.custom_query(query_text)
+    return response.response
+
+# Define Gradio interface
+md_json_objs = []
+
+def upload_wrapper(uploaded_file):
+    global md_json_objs
+    md_json_objs = upload_and_process_file(uploaded_file)
+    return "File successfully processed!"
+
+iface = gr.Interface(
+    fn=ask_question,
+    inputs=[
+        gr.inputs.State(),
+        gr.inputs.Textbox(label="Enter your query:"),
+        gr.inputs.File(label="Upload a query image (if any):", optional=True)
+    ],
+    outputs="text",
+    title="Medical Knowledge Base & Query System"
+)
+
+upload_iface = gr.Interface(
+    fn=upload_wrapper,
+    inputs=gr.inputs.File(label="Upload a medical textbook (pdf):"),
+    outputs="text",
+    title="Upload Knowledge Base"
+)
+
+app = gr.TabbedInterface([upload_iface, iface], ["Upload Knowledge Base", "Ask a Question"])
+app.launch()