Spaces:
Running
Running
File size: 6,436 Bytes
aea1ea5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 |
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") |