| import os |
| import faiss |
| import pickle |
| import numpy as np |
| from typing import List, Dict |
| from docarray import Document, DocumentArray |
| from jina import Executor, requests |
| from sentence_transformers import SentenceTransformer |
| from transformers import BlipProcessor, BlipForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM |
| from pdfminer.high_level import extract_text |
| import fitz |
| from PIL import Image |
| import traceback |
| import torch |
| import re |
|
|
|
|
| class MultimodalRAGExecutor(Executor): |
| def __init__( |
| self, |
| llm_model_name: str = "google/flan-t5-large", |
| embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2", |
| image_model: str = "Salesforce/blip-image-captioning-base", |
| index_file: str = "faiss_index.bin", |
| metadata_file: str = "metadata.pkl", |
| dim: int = 384, |
| use_huggingface_api: bool = False, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
|
|
| self.llm_model_name = llm_model_name |
| self.embedding_model = embedding_model |
| self.image_model = image_model |
| self.index_file = index_file |
| self.metadata_file = metadata_file |
| self.dim = dim |
| self.use_huggingface_api = use_huggingface_api |
|
|
| |
| print(f"📥 Loading embedding model: {embedding_model}...") |
| self.embedder = SentenceTransformer(embedding_model) |
| print("✅ Embedding model loaded!") |
|
|
| |
| print(f"📥 Loading image model: {image_model}...") |
| try: |
| self.blip_processor = BlipProcessor.from_pretrained(image_model) |
| self.blip_model = BlipForConditionalGeneration.from_pretrained(image_model) |
| print("✅ Image model loaded!") |
| except Exception as e: |
| print(f"⚠️ Warning: Could not load BLIP model: {e}") |
| self.blip_processor = None |
| self.blip_model = None |
|
|
| |
| self._initialize_llm() |
|
|
| |
| self._load_or_create_index() |
|
|
| def _initialize_llm(self): |
| """Initialize LLM (API or local)""" |
| if self.use_huggingface_api: |
| hf_token = os.getenv("HUGGINGFACE_API_KEY") |
| if hf_token: |
| try: |
| from huggingface_hub import InferenceClient |
| self.llm_client = InferenceClient(token=hf_token) |
| self.llm_mode = "api" |
| print("✅ Using HuggingFace API for LLM!") |
| return |
| except Exception as e: |
| print(f"⚠️ HF API initialization failed: {e}") |
| |
| |
| print(f"📥 Loading local LLM: {self.llm_model_name}...") |
| try: |
| self.tokenizer = AutoTokenizer.from_pretrained(self.llm_model_name) |
| self.llm_model = AutoModelForSeq2SeqLM.from_pretrained(self.llm_model_name) |
| self.llm_mode = "local" |
| print("✅ Local LLM loaded!") |
| except Exception as e: |
| print(f"⚠️ LLM loading failed: {e}") |
| self.llm_mode = "extractive" |
| print("⚠️ Falling back to extractive QA") |
|
|
| def _load_or_create_index(self): |
| """Load existing index or create new one""" |
| if os.path.exists(self.index_file) and os.path.exists(self.metadata_file): |
| try: |
| self.index = faiss.read_index(self.index_file) |
| with open(self.metadata_file, "rb") as f: |
| self.metadata = pickle.load(f) |
| print(f"✅ Loaded existing index with {self.index.ntotal} vectors") |
| except Exception as e: |
| print(f"⚠️ Could not load index: {e}. Creating new one.") |
| self.index = faiss.IndexFlatL2(self.dim) |
| self.metadata = [] |
| else: |
| self.index = faiss.IndexFlatL2(self.dim) |
| self.metadata = [] |
| print("✅ Created new empty index") |
|
|
| def _get_embedding(self, text: str) -> np.ndarray: |
| """Generate embeddings using sentence-transformers""" |
| embedding = self.embedder.encode(text, convert_to_numpy=True) |
| return embedding.astype(np.float32) |
|
|
| def _analyze_image(self, image: Image.Image, context: str = "") -> str: |
| """Analyze image using BLIP model""" |
| if not self.blip_processor or not self.blip_model: |
| return "Image analysis unavailable" |
| |
| try: |
| |
| inputs = self.blip_processor(image, return_tensors="pt") |
| out = self.blip_model.generate(**inputs, max_length=100) |
| caption = self.blip_processor.decode(out[0], skip_special_tokens=True) |
| |
| |
| text_prompt = "a detailed description of" |
| inputs = self.blip_processor(image, text_prompt, return_tensors="pt") |
| out = self.blip_model.generate(**inputs, max_length=150) |
| detailed = self.blip_processor.decode(out[0], skip_special_tokens=True) |
| |
| description = f"Caption: {caption}. Details: {detailed}" |
| if context: |
| description += f". Context: {context[:100]}" |
| |
| return description |
| except Exception as e: |
| print(f"Image analysis error: {e}") |
| return "Image analysis failed" |
|
|
| def _extract_text_from_pdf(self, pdf_path: str) -> List[str]: |
| """Extract text from PDF using PyMuPDF""" |
| texts = [] |
| try: |
| doc = fitz.open(pdf_path) |
| for page_num, page in enumerate(doc, start=1): |
| text = page.get_text("text") |
| if text and text.strip(): |
| texts.append(f"Page {page_num}:\n{text.strip()}") |
| doc.close() |
| except Exception as e: |
| print(f"Text extraction error: {e}") |
| return texts |
|
|
| def _extract_images_from_pdf(self, pdf_path: str) -> List[Dict]: |
| """Extract images from PDF using PyMuPDF""" |
| images_data = [] |
| try: |
| doc = fitz.open(pdf_path) |
| for page_num, page in enumerate(doc, start=1): |
| image_list = page.get_images(full=True) |
| |
| for img_idx, img in enumerate(image_list): |
| try: |
| xref = img[0] |
| base_image = doc.extract_image(xref) |
| image_bytes = base_image["image"] |
| |
| |
| import io |
| pil_image = Image.open(io.BytesIO(image_bytes)).convert('RGB') |
| |
| |
| width, height = pil_image.size |
| if width >= 100 and height >= 100: |
| images_data.append({ |
| 'image': pil_image, |
| 'page': page_num, |
| 'index': img_idx |
| }) |
| except Exception as e: |
| print(f"Error extracting image on page {page_num}: {e}") |
| continue |
| |
| doc.close() |
| except Exception as e: |
| print(f"Image extraction error: {e}") |
| |
| return images_data |
|
|
| def _generate_answer_local(self, prompt: str) -> str: |
| """Generate answer using local FLAN-T5 model""" |
| try: |
| inputs = self.tokenizer( |
| prompt, |
| return_tensors="pt", |
| max_length=512, |
| truncation=True |
| ) |
| |
| outputs = self.llm_model.generate( |
| inputs.input_ids, |
| max_new_tokens=200, |
| temperature=0.3, |
| do_sample=False, |
| num_beams=4 |
| ) |
| |
| answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True) |
| return answer.strip() |
| except Exception as e: |
| print(f"Local generation error: {e}") |
| return "" |
|
|
| def _generate_answer_api(self, prompt: str) -> str: |
| """Generate answer using HuggingFace API""" |
| try: |
| response = self.llm_client.text_generation( |
| prompt, |
| model="google/flan-t5-xxl", |
| max_new_tokens=250, |
| temperature=0.3 |
| ) |
| return response.strip() |
| except Exception as e: |
| print(f"API generation error: {e}") |
| return "" |
|
|
| def _generate_extractive_answer(self, query: str, context: str) -> str: |
| """Fallback extractive QA""" |
| sentences = re.split(r'[.!?]+', context) |
| sentences = [s.strip() for s in sentences if len(s.strip()) > 20] |
| |
| query_words = set(query.lower().split()) |
| scored = [] |
| |
| for sent in sentences[:30]: |
| sent_words = set(sent.lower().split()) |
| overlap = len(query_words.intersection(sent_words)) |
| if overlap > 0: |
| scored.append((overlap, sent)) |
| |
| scored.sort(reverse=True, key=lambda x: x[0]) |
| |
| if scored: |
| top_sentences = [s for _, s in scored[:3]] |
| return ". ".join(top_sentences) + "." |
| |
| return "I couldn't find relevant information to answer your question." |
|
|
| def _generate_answer(self, prompt: str, query: str = "", context: str = "") -> str: |
| """Generate answer with fallback chain""" |
| if self.llm_mode == "api": |
| answer = self._generate_answer_api(prompt) |
| if answer and len(answer) > 20: |
| return answer |
| |
| if self.llm_mode == "local": |
| answer = self._generate_answer_local(prompt) |
| if answer and len(answer) > 20: |
| return answer |
| |
| |
| return self._generate_extractive_answer(query, context) |
|
|
| def _clean_answer(self, answer: str) -> str: |
| """Clean unwanted artifacts from answer""" |
| answer = re.sub(r'^(Answer|Context|Question):\s*', '', answer, flags=re.IGNORECASE) |
| answer = re.sub(r'[#*_`]+', '', answer) |
| answer = re.sub(r'\s+', ' ', answer).strip() |
| return answer |
|
|
| @requests(on="/upload") |
| def upload(self, docs: DocumentArray, **kwargs): |
| """Upload PDF, extract text + images, index them""" |
| try: |
| for doc in docs: |
| file_path = doc.uri |
| if file_path.startswith("file://"): |
| file_path = file_path.replace("file://", "") |
|
|
| if not os.path.exists(file_path) or not file_path.endswith(".pdf"): |
| continue |
|
|
| print(f"📄 Processing PDF: {file_path}") |
|
|
| |
| text_chunks = self._extract_text_from_pdf(file_path) |
| |
| |
| text_count = 0 |
| context_text = "\n".join(text_chunks[:2])[:1000] |
| |
| for chunk in text_chunks: |
| |
| paragraphs = chunk.split("\n\n") |
| for para in paragraphs: |
| if para.strip() and len(para.strip()) > 50: |
| emb = self._get_embedding(para.strip()) |
| self.index.add(np.array([emb])) |
| self.metadata.append({ |
| "type": "text", |
| "content": para.strip() |
| }) |
| text_count += 1 |
|
|
| print(f"✅ Indexed {text_count} text chunks") |
|
|
| |
| images_data = self._extract_images_from_pdf(file_path) |
| image_count = 0 |
| |
| for img_data in images_data: |
| |
| description = self._analyze_image(img_data['image'], context_text) |
| |
| |
| img_path = os.path.abspath( |
| f"img_p{img_data['page']}_i{img_data['index']}.png" |
| ) |
| img_data['image'].save(img_path, "PNG") |
| |
| |
| embed_text = f"Image from page {img_data['page']}: {description}" |
| emb = self._get_embedding(embed_text) |
| self.index.add(np.array([emb])) |
| self.metadata.append({ |
| "type": "image", |
| "content": f"file://{img_path}", |
| "description": description, |
| "page": img_data['page'] |
| }) |
| image_count += 1 |
|
|
| print(f"✅ Analyzed and indexed {image_count} images") |
|
|
| |
| faiss.write_index(self.index, self.index_file) |
| with open(self.metadata_file, "wb") as f: |
| pickle.dump(self.metadata, f) |
|
|
| summary = f"✅ Upload complete!\n" |
| summary += f"📊 Total vectors: {self.index.ntotal}\n" |
| summary += f"📝 Text chunks: {text_count}\n" |
| summary += f"🖼️ Images: {image_count}\n" |
| summary += f"💚 100% FREE processing!" |
| |
| return DocumentArray([Document(text=summary)]) |
| |
| except Exception as e: |
| error_msg = f"❌ Upload failed:\n{traceback.format_exc()}" |
| print(error_msg) |
| return DocumentArray([Document(text=error_msg)]) |
|
|
| @requests(on="/query") |
| def query(self, docs: DocumentArray, **kwargs): |
| """Query with multimodal RAG""" |
| results = DocumentArray() |
|
|
| if self.index.ntotal == 0: |
| return DocumentArray([ |
| Document(text="⚠️ No documents uploaded. Please upload a PDF first.") |
| ]) |
|
|
| for doc in docs: |
| try: |
| query_text = doc.text |
| |
| |
| query_emb = self._get_embedding(query_text) |
| D, I = self.index.search(np.array([query_emb]), k=8) |
|
|
| context_parts = [] |
| matched_images = [] |
| image_descriptions = [] |
| |
| for idx in I[0]: |
| if idx < len(self.metadata): |
| meta = self.metadata[idx] |
| if meta["type"] == "text": |
| context_parts.append(meta["content"]) |
| elif meta["type"] == "image": |
| matched_images.append(Document(uri=meta["content"])) |
| image_descriptions.append( |
| f"[Image Page {meta.get('page', '?')}]: {meta['description']}" |
| ) |
|
|
| |
| context_text = "\n\n".join(context_parts[:4]) |
| if image_descriptions: |
| context_text += "\n\nRelevant Images:\n" + "\n".join(image_descriptions[:2]) |
|
|
| |
| if len(context_text) > 2000: |
| context_text = context_text[:2000] + "..." |
|
|
| |
| prompt = ( |
| f"Answer the question based on the context from a research paper.\n\n" |
| f"Context:\n{context_text}\n\n" |
| f"Question: {query_text}\n\n" |
| f"Answer:" |
| ) |
|
|
| |
| answer = self._generate_answer(prompt, query_text, context_text) |
| answer = self._clean_answer(answer) |
|
|
| |
| answer_doc = Document(text=answer) |
| if matched_images: |
| answer_doc.chunks = DocumentArray(matched_images[:4]) |
|
|
| results.append(answer_doc) |
| |
| except Exception as e: |
| error_msg = f"❌ Query failed: {str(e)}\n{traceback.format_exc()}" |
| print(error_msg) |
| results.append(Document(text=error_msg)) |
| |
| return results |
|
|
| @requests(on="/stats") |
| def stats(self, docs: DocumentArray, **kwargs): |
| """Get index statistics""" |
| text_count = sum(1 for m in self.metadata if m["type"] == "text") |
| image_count = sum(1 for m in self.metadata if m["type"] == "image") |
| |
| stats_text = ( |
| f"📊 Index Statistics:\n" |
| f"Total vectors: {self.index.ntotal}\n" |
| f"Text chunks: {text_count}\n" |
| f"Images: {image_count}\n" |
| f"LLM mode: {self.llm_mode}\n" |
| f"💚 100% FREE models!" |
| ) |
| |
| return DocumentArray([Document(text=stats_text)]) |
|
|
| @requests(on="/reset") |
| def reset(self, docs: DocumentArray, **kwargs): |
| """Reset index and metadata""" |
| try: |
| self.index = faiss.IndexFlatL2(self.dim) |
| self.metadata = [] |
| |
| if os.path.exists(self.index_file): |
| os.remove(self.index_file) |
| if os.path.exists(self.metadata_file): |
| os.remove(self.metadata_file) |
| |
| return DocumentArray([Document(text="✅ Index reset successfully!")]) |
| except Exception as e: |
| return DocumentArray([ |
| Document(text=f"❌ Reset failed: {str(e)}") |
| ]) |
