Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| import torch | |
| 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 | |
| class MultimodalRAG: | |
| 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.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} | |
| ) | |
| 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 process_image(self, image_path): | |
| if not os.path.exists(image_path): | |
| print(f"Warning: Image path {image_path} does not exist") | |
| return None | |
| image = Image.open(image_path) | |
| 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 | |
| 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: | |
| docs = self.retrieve_related_documents(query_text, image_path) | |
| else: | |
| docs = self.retrieve_related_documents(query_text) | |
| result = self.qa_chain({"query": query_text}) | |
| answer = result["result"] | |
| sources = [doc.page_content[:1000] + "..." for doc in result["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", [] | |
| if pdf_file is None: | |
| return "Please upload a PDF document first", [] | |
| 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 __name__ == "__main__": | |
| demo.launch(share=True, server_name="0.0.0.0") |