Update app.py

app.py

@@ -1,184 +1,100 @@
-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=None):
+    def __init__(self, pdf_path):
         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.pdf_path = pdf_path
-        self.documents = []
-        self.vector_store = None
-        self.retriever = None
-        self.qa_chain = None
-        
+
+        self.documents = self._load_and_split(pdf_path)
+        self.vector_store = FAISS.from_documents(self.documents, self.text_embeddings)
+
         try:
             self.llm = HuggingFacePipeline.from_model_id(
                 model_id="google/flan-t5-large",
                 task="text2text-generation",
-                model_kwargs={"temperature": 0.7, "max_length": 512}
+                model_kwargs={"temperature": 0.7, "max_length": 512, "device": -1}
             )
-        except Exception as e:
-            print(f"Error loading flan-t5 model: {e}")
+        except Exception:
             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 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 _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)
 
-        image = Image.open(image_path)
+    def _get_image_features(self, image_path):
+        image = Image.open(image_path).convert("RGB")
         inputs = self.processor(images=image, return_tensors="pt")
         with torch.no_grad():
-            image_features = self.vision_model.get_image_features(**inputs)
-        return image_features
-
-    def generate_image_description(self, image_features):
-        return "a photo"
-
-    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
+            return self.vision_model.get_image_features(**inputs)
 
-        docs = self.retriever.get_relevant_documents(enhanced_query)
-        return docs
+    def _generate_image_description(self, image_features):
+        return "an image"
 
     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:
-            docs = self.retrieve_related_documents(query_text, image_path)
+            feats = self._get_image_features(image_path)
+            img_desc = self._generate_image_description(feats)
+            full_query = f"{query_text} {img_desc}"
         else:
-            docs = self.retrieve_related_documents(query_text)
+            full_query = query_text
 
-        result = self.qa_chain({"query": query_text})
-        
+        result = self.qa_chain({"query": full_query})
         answer = result["result"]
-        sources = [doc.page_content[:1000] + "..." for doc in result["source_documents"]]
-        
+        sources = [doc.metadata for doc in result.get("source_documents", [])]
         return answer, sources
 
-rag_system = MultimodalRAG()
 
-def upload_pdf(pdf_file):
-    if pdf_file is None:
-        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", []
-    
+def run_rag(pdf_file, query, image_file=None):
     if pdf_file is None:
-        return "Please upload a PDF document first", []
-    
+        return "Please upload a PDF.", []
+
+    pdf_path = pdf_file.name
     image_path = None
-    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)}", []
-
-# Create Gradio interface
-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 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 __name__ == "__main__":
-    demo.launch(share=True, server_name="0.0.0.0")
+    iface.launch(server_name="0.0.0.0", server_port=7860)