functional

app.py

@@ -1,3 +1,5 @@
+# TODO: return all pages used to form answer
+
 import torch
 from transformers import BitsAndBytesConfig
 from llama_index.llms.huggingface import HuggingFaceLLM
@@ -6,7 +8,7 @@ from llama_index.core import SimpleDirectoryReader
 from llama_index.core import VectorStoreIndex, SummaryIndex
 from llama_index.core.prompts import PromptTemplate
 from llama_index.core import Settings
-
+from PIL import Image
 
 import gradio as gr
 
@@ -63,6 +65,9 @@ def load_RAG_pipeline():
     # raw data
     documents = SimpleDirectoryReader("assets/txts").load_data()
     vector_index = VectorStoreIndex.from_documents(documents)
+    # vector_index.persist(persist_dir="vectors")
+    # https://docs.llamaindex.ai/en/v0.10.17/understanding/storing/storing.html
+
     # summary_index = SummaryIndex.from_documents(documents)
     query_engine = vector_index.as_query_engine(response_mode="compact", similarity_top_k=3)
     return query_engine
@@ -78,23 +83,31 @@ def get_answer(question, temperature, nucleus_sampling, max_tokens):
     # For example, you could use a machine learning model for RAG.
     # answer = "This is a placeholder answer."
     # https://docs.llamaindex.ai/en/stable/module_guides/supporting_modules/settings/#setting-local-configurations
-    return query_engine.query(question)
+    response = query_engine.query(question)
+    return response
 
 
-def get_answer_page(question):
+def get_answer_page(response):
     # Implement logic to retrieve the page number or an image of the page with the answer.
-    answer_page = "Page X - placeholder image."
-    return answer_page
+    # best image
+    best_match = response.source_nodes[0].metadata["file_path"]
+    answer_page = float(int(best_match[-8:-4]))
+    image = Image.open(best_match.replace("txt", "png"))
+    return image, answer_page
 
 
 # Create the gr.Interface function
 def ask_my_thesis(question, temperature, nucleus_sampling, max_tokens):
     answer = get_answer(question, temperature, nucleus_sampling, max_tokens)
-    answer_page = get_answer_page(question)
-    return answer, answer_page
+    image, answer_page = get_answer_page(answer)
+    return answer, image, answer_page
 
 
 # Set up the interface options based on the design in the image.
+output_image = gr.Image(label="Answer Page")
+
+# examples
+
 iface = gr.Interface(
     fn=ask_my_thesis,
     inputs=[
@@ -103,12 +116,17 @@ iface = gr.Interface(
         gr.Slider(0, 1, value=0.9, label="Nucleus Sampling"),
         gr.Slider(1, 500, value=100, label="Max Generated Number of Tokens"),
     ],
-    outputs=[gr.Textbox(label="Answer"), gr.Image(label="Answer Page")],
-    title="Ask my thesis",
-    description="Chat with the manuscript: ask questions and receive answers with references.",
+    outputs=[gr.Textbox(label="Answer"), output_image, gr.Label()],
+    title="Ask my thesis: Intelligent Automation for AI-Driven Document Understanding",
+    description=r"""Chat with the thesis manuscript: ask questions and receive answers with multimodal references (WIP).
+    
+    Spoiler: RAG application with LLM and embedding vector store can be quite slow on a 290 page document ;D
+    """,
     allow_flagging="never",
 )
+# https://github.com/gradio-app/gradio/issues/4309
 
+# https://discuss.huggingface.co/t/add-background-image/16381/4 background image
 # Start the application.
 if __name__ == "__main__":
     iface.launch()