Update app.py

app.py

@@ -1,100 +1,184 @@
+import gradio as gr
 import os
 import torch
-from PIL import Image
-import gradio as gr
 from transformers import AutoProcessor, AutoModel
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain.chains import RetrievalQA
 from langchain_community.llms import HuggingFacePipeline
+from PIL import Image
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
-os.environ["CUDA_VISIBLE_DEVICES"] = ""
-
 class MultimodalRAG:
-    def __init__(self, pdf_path):
+    def __init__(self, pdf_path=None):
         self.processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32")
         self.vision_model = AutoModel.from_pretrained("openai/clip-vit-base-patch32")
         self.text_embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-
-        self.documents = self._load_and_split(pdf_path)
-        self.vector_store = FAISS.from_documents(self.documents, self.text_embeddings)
-
+        self.pdf_path = pdf_path
+        self.documents = []
+        self.vector_store = None
+        self.retriever = None
+        self.qa_chain = None
+        
         try:
             self.llm = HuggingFacePipeline.from_model_id(
                 model_id="google/flan-t5-large",
                 task="text2text-generation",
-                model_kwargs={"temperature": 0.7, "max_length": 512, "device": -1}
+                model_kwargs={"temperature": 0.7, "max_length": 512}
             )
-        except Exception:
+        except Exception as e:
+            print(f"Error loading flan-t5 model: {e}")
             from langchain.llms import OpenAI
             self.llm = OpenAI(temperature=0.7)
+            
+        if pdf_path and os.path.exists(pdf_path):
+            self.load_pdf(pdf_path)
+
+    def load_pdf(self, pdf_path):
+        if not os.path.exists(pdf_path):
+            raise FileNotFoundError(f"PDF file not found: {pdf_path}")
+
+        loader = PyPDFLoader(pdf_path)
+        self.documents = loader.load()
 
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=1000,
+            chunk_overlap=200
+        )
+        self.documents = text_splitter.split_documents(self.documents)
+        
+        self.vector_store = FAISS.from_documents(self.documents, self.text_embeddings)
+        
         self.retriever = self.vector_store.as_retriever(search_kwargs={"k": 2})
+        
         self.qa_chain = RetrievalQA.from_chain_type(
             llm=self.llm,
             chain_type="stuff",
             retriever=self.retriever,
             return_source_documents=True
         )
+        
+        return f"Successfully loaded and processed PDF: {pdf_path}"
 
-    def _load_and_split(self, pdf_path):
-        loader = PyPDFLoader(pdf_path)
-        docs = loader.load()
-        splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-        return splitter.split_documents(docs)
+    def process_image(self, image_path):
+        if not os.path.exists(image_path):
+            print(f"Warning: Image path {image_path} does not exist")
+            return None
 
-    def _get_image_features(self, image_path):
-        image = Image.open(image_path).convert("RGB")
+        image = Image.open(image_path)
         inputs = self.processor(images=image, return_tensors="pt")
         with torch.no_grad():
-            return self.vision_model.get_image_features(**inputs)
+            image_features = self.vision_model.get_image_features(**inputs)
+        return image_features
 
-    def _generate_image_description(self, image_features):
+    def generate_image_description(self, image_features):
         return "an image"
 
+    def retrieve_related_documents(self, query_text, image_path=None):
+        if image_path:
+            image_features = self.process_image(image_path)
+
+            if image_features is not None:
+                image_query = self.generate_image_description(image_features)
+
+                enhanced_query = f"{query_text} {image_query}"
+            else:
+                enhanced_query = query_text
+        else:
+            enhanced_query = query_text
+
+        docs = self.retriever.get_relevant_documents(enhanced_query)
+        return docs
+
     def answer_query(self, query_text, image_path=None):
+        if not self.vector_store or not self.qa_chain:
+            return "Please upload a PDF document first.", []
+            
         if image_path:
-            feats = self._get_image_features(image_path)
-            img_desc = self._generate_image_description(feats)
-            full_query = f"{query_text} {img_desc}"
+            docs = self.retrieve_related_documents(query_text, image_path)
         else:
-            full_query = query_text
+            docs = self.retrieve_related_documents(query_text)
 
-        result = self.qa_chain({"query": full_query})
+        result = self.qa_chain({"query": query_text})
+        
         answer = result["result"]
-        sources = [doc.metadata for doc in result.get("source_documents", [])]
+        sources = [doc.page_content[:1000] + "..." for doc in result["source_documents"]]
+        
         return answer, sources
 
+rag_system = MultimodalRAG()
 
-def run_rag(pdf_file, query, image_file=None):
+def upload_pdf(pdf_file):
     if pdf_file is None:
-        return "Please upload a PDF.", []
-
-    pdf_path = pdf_file.name
+        return "No file uploaded"
+    
+    file_path = pdf_file.name
+    try:
+        result = rag_system.load_pdf(file_path)
+        return result
+    except Exception as e:
+        return f"Error processing PDF: {str(e)}"
+
+def save_image(image):
+    if image is None:
+        return None
+    
+    temp_path = "temp_image.jpg"
+    image.save(temp_path)
+    return temp_path
+
+def process_query(query, pdf_file, image=None):
+    if not query.strip():
+        return "Please enter a question", []
+    
+    if pdf_file is None:
+        return "Please upload a PDF document first", []
+    
     image_path = None
-    if image_file:
-        image_path = image_file.name
-
-    rag = MultimodalRAG(pdf_path)
-    answer, sources = rag.answer_query(query, image_path)
-    return answer, sources
-
-iface = gr.Interface(
-    fn=run_rag,
-    inputs=[
-        gr.File(label="PDF Document", file_types=[".pdf"]),
-        gr.Textbox(label="Query", placeholder="Enter your question here..."),
-        gr.File(label="Optional Image", file_types=[".png", ".jpg", ".jpeg"], optional=True)
-    ],
-    outputs=[
-        gr.Textbox(label="Answer"),
-        gr.JSON(label="Source Documents")
-    ],
-    title="Multimodal RAG QA",
-    description="Upload a PDF, ask a question, optionally provide an image."
-)
+    if image is not None:
+        image_path = save_image(image)
+    
+    try:
+        answer, sources = rag_system.answer_query(query, image_path)
+        if image_path and os.path.exists(image_path):
+            os.remove(image_path)
+        return answer, sources
+    except Exception as e:
+        if image_path and os.path.exists(image_path):
+            os.remove(image_path)
+        return f"Error processing query: {str(e)}", []
+
+
+with gr.Blocks(title="Multimodal RAG System") as demo:
+    gr.Markdown("# Multimodal RAG System")
+    gr.Markdown("Upload a PDF document and ask questions about it. You can also add an image for multimodal context.")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            pdf_input = gr.File(label="Upload PDF Document")
+            upload_button = gr.Button("Process PDF")
+            status_output = gr.Textbox(label="Status")
+            
+            upload_button.click(
+                fn=upload_pdf,
+                inputs=[pdf_input],
+                outputs=[status_output]
+            )
+        
+        with gr.Column(scale=2):
+            image_input = gr.Image(label="Optional: Upload an Image", type="pil")
+            query_input = gr.Textbox(label="Ask a question")
+            submit_button = gr.Button("Submit Question")
+            
+            answer_output = gr.Textbox(label="Answer")
+            sources_output = gr.JSON(label="Sources")
+            
+            submit_button.click(
+                fn=process_query,
+                inputs=[query_input, pdf_input, image_input],
+                outputs=[answer_output, sources_output]
+            )
 
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(share=True, server_name="0.0.0.0")