Spaces:
Sleeping
Sleeping
| # Main Multimodal RAG System | |
| import os | |
| import json | |
| import uuid | |
| import numpy as np | |
| import faiss | |
| import torch | |
| from datetime import datetime | |
| from pathlib import Path | |
| from typing import List, Dict | |
| import streamlit as st | |
| from config import (DATA_DIRECTORY, DATA_FILES, DEFAULT_TOP_K, MAX_CONTEXT_LENGTH, LLM_TEMPERATURE, FAISS_TEMP_MEMORY) | |
| from models import load_models | |
| from document_processor import DocumentProcessor | |
| from device_setup import clear_gpu_memory | |
| class MultimodalRAG: | |
| """Main RAG system that handles document storage, search, and generation""" | |
| def __init__(self): | |
| # load all models | |
| model_components = load_models() | |
| if model_components[0] is None: | |
| raise Exception("Failed to load models") | |
| (self.text_encoder, self.clip_model, self.clip_processor, self.whisper_model, self.ocr_reader, self.llm, self.device, self.has_gpu) = model_components | |
| # initialize document processor | |
| self.processor = DocumentProcessor( | |
| self.clip_model, self.clip_processor, self.whisper_model, | |
| self.ocr_reader, self.text_encoder, self.device | |
| ) | |
| # initialize storage | |
| self.documents = [] | |
| self.metadata = [] | |
| self.embeddings = None | |
| self.index = None | |
| # setup data directory | |
| self.data_dir = DATA_DIRECTORY | |
| self.data_dir.mkdir(exist_ok=True) | |
| # load existing data | |
| self.load_existing_data() | |
| def clear_knowledge_base(self): | |
| """clear all stored documents and embeddings""" | |
| self.documents = [] | |
| self.metadata = [] | |
| self.embeddings = None | |
| self.index = None | |
| try: | |
| # remove all files in data directory | |
| for file in self.data_dir.glob("*"): | |
| file.unlink() | |
| except Exception as e: | |
| st.error(f"Error clearing saved data files: {e}") | |
| def save_data(self): | |
| """Save documents, metadata, embeddings, and FAISS index to disk""" | |
| try: | |
| # save documents and metadata as JSON | |
| with open(self.data_dir / DATA_FILES["documents"], "w") as f: | |
| json.dump(self.documents, f, default=str) | |
| with open(self.data_dir / DATA_FILES["metadata"], "w") as f: | |
| json.dump(self.metadata, f, default=str) | |
| # save FAISS index | |
| if self.index is not None: | |
| if self.has_gpu: | |
| # convert GPU index back to CPU for saving | |
| cpu_index = faiss.index_gpu_to_cpu(self.index) | |
| faiss.write_index(cpu_index, str(self.data_dir / DATA_FILES["faiss_index"])) | |
| else: | |
| faiss.write_index(self.index, str(self.data_dir / DATA_FILES["faiss_index"])) | |
| # save embeddings | |
| if self.embeddings is not None: | |
| np.save(self.data_dir / DATA_FILES["embeddings"], self.embeddings) | |
| except Exception as e: | |
| st.error(f"Error saving data: {e}") | |
| def load_existing_data(self): | |
| """Load existing documents and metadata from disk""" | |
| try: | |
| # load documents | |
| doc_file = self.data_dir / DATA_FILES["documents"] | |
| if doc_file.exists(): | |
| with open(doc_file, "r") as f: | |
| self.documents = json.load(f) | |
| # load metadata | |
| meta_file = self.data_dir / DATA_FILES["metadata"] | |
| if meta_file.exists(): | |
| with open(meta_file, "r") as f: | |
| self.metadata = json.load(f) | |
| # load FAISS index | |
| index_file = self.data_dir / DATA_FILES["faiss_index"] | |
| if index_file.exists(): | |
| if self.has_gpu: | |
| cpu_index = faiss.read_index(str(index_file)) | |
| res = faiss.StandardGpuResources() | |
| self.index = faiss.index_cpu_to_gpu(res, 0, cpu_index) | |
| else: | |
| self.index = faiss.read_index(str(index_file)) | |
| # load embeddings | |
| emb_file = self.data_dir / DATA_FILES["embeddings"] | |
| if emb_file.exists(): | |
| self.embeddings = np.load(emb_file) | |
| except Exception as e: | |
| st.error(f"Error loading existing data: {e}") | |
| def process_file(self, uploaded_file, file_type: str) -> bool: | |
| """Process uploaded file and add to knowledge base""" | |
| try: | |
| # process the file | |
| processed_data = self.processor.process_uploaded_file(uploaded_file, file_type, self.data_dir) | |
| if processed_data is None: | |
| return False | |
| # create document record | |
| file_id = str(uuid.uuid4()) | |
| timestamp = datetime.now().isoformat() | |
| document = { | |
| 'id': file_id, | |
| 'type': file_type if file_type not in ['png', 'jpg', 'jpeg'] else 'image', | |
| 'filename': uploaded_file.name, | |
| 'content': processed_data['content'], | |
| 'timestamp': timestamp, | |
| 'embedding_type': processed_data['embedding_type'] | |
| } | |
| # handle file-specific storage (images and audio) | |
| if file_type in ['png', 'jpg', 'jpeg', 'mp3', 'wav', 'ogg', 'm4a']: | |
| # save the actual file | |
| file_save_path = self.data_dir / f"{file_id}.{file_type}" | |
| uploaded_file.seek(0) # Reset file pointer | |
| with open(file_save_path, 'wb') as f: | |
| f.write(uploaded_file.read()) | |
| document[f"{document['type']}_path"] = str(file_save_path) | |
| # add embedding to collection | |
| embedding = processed_data['embedding'] | |
| if self.embeddings is None: | |
| self.embeddings = embedding.reshape(1, -1) | |
| else: | |
| # handle dimension mismatch (for mixed embedding types) | |
| current_dim = self.embeddings.shape[1] | |
| if embedding.shape[0] != current_dim: | |
| if embedding.shape[0] > current_dim: | |
| embedding = embedding[:current_dim] | |
| else: | |
| padding = np.zeros(current_dim - embedding.shape[0]) | |
| embedding = np.concatenate([embedding, padding]) | |
| self.embeddings = np.vstack([self.embeddings, embedding.reshape(1, -1)]) | |
| # Add to collections | |
| self.documents.append(document) | |
| self.metadata.append({ | |
| 'id': file_id, | |
| 'type': document['type'], | |
| 'filename': uploaded_file.name, | |
| 'timestamp': timestamp, | |
| 'content_preview': processed_data['content'][:200] + "..." if len(processed_data['content']) > 200 else processed_data['content'] | |
| }) | |
| # Rebuild search index | |
| self.build_faiss_index() | |
| self.save_data() | |
| clear_gpu_memory() | |
| return True | |
| except Exception as e: | |
| st.error(f"Error processing file: {e}") | |
| return False | |
| def build_faiss_index(self): | |
| """Build or rebuild FAISS index for similarity search""" | |
| try: | |
| if self.embeddings is not None and len(self.embeddings) > 0: | |
| dimension = self.embeddings.shape[1] | |
| if self.has_gpu: | |
| try: | |
| # Create GPU resources with limited temp memory | |
| res = faiss.StandardGpuResources() | |
| res.setTempMemory(FAISS_TEMP_MEMORY) | |
| cpu_index = faiss.IndexFlatIP(dimension) | |
| self.index = faiss.index_cpu_to_gpu(res, 0, cpu_index) | |
| except RuntimeError as e: | |
| st.warning(f"GPU FAISS failed: {e}. Using CPU indexing.") | |
| self.index = faiss.IndexFlatIP(dimension) | |
| else: | |
| self.index = faiss.IndexFlatIP(dimension) | |
| # Normalize embeddings and add to index | |
| embeddings_normalized = self.embeddings.copy().astype('float32') | |
| faiss.normalize_L2(embeddings_normalized) | |
| self.index.add(embeddings_normalized) | |
| except Exception as e: | |
| st.error(f"Error building FAISS index: {e}") | |
| def search(self, query: str, top_k: int = DEFAULT_TOP_K) -> List[Dict]: | |
| """Search for relevant documents""" | |
| try: | |
| if self.index is None or len(self.documents) == 0: | |
| return [] | |
| # get query embedding | |
| query_embedding = self.processor.get_text_embedding(query) | |
| query_embedding = query_embedding.reshape(1, -1).astype('float32') | |
| faiss.normalize_L2(query_embedding) | |
| # search | |
| scores, indices = self.index.search(query_embedding, min(top_k, len(self.documents))) | |
| # build results | |
| results = [] | |
| for score, idx in zip(scores[0], indices[0]): | |
| if idx < len(self.documents): | |
| result = self.documents[idx].copy() | |
| result['score'] = float(score) | |
| results.append(result) | |
| return results | |
| except Exception as e: | |
| st.error(f"Error searching: {e}") | |
| return [] | |
| def generate_answer(self, query: str, retrieved_docs: List[Dict]) -> str: | |
| """Generate answer using local LLM""" | |
| try: | |
| if not retrieved_docs: | |
| return "No relevant documents found to answer your query." | |
| # build context from retrieved documents | |
| context = "Context from retrieved documents:\n\n" | |
| citations = [] | |
| for i, doc in enumerate(retrieved_docs[:3]): # using top 3 documents | |
| context += f"[{i+1}] {doc['filename']} ({doc['type']}): {doc['content'][:MAX_CONTEXT_LENGTH]}...\n\n" | |
| citations.append(f"[{i+1}] {doc['filename']}") | |
| # Create prompt | |
| prompt = f"""Based on the following context, answer the user's question clearly and concisely. Use citations [1], [2], etc. | |
| Context: | |
| {context} | |
| Question: {query} | |
| Answer:""" | |
| # generate response | |
| response = self.llm.generate(prompt, temp=LLM_TEMPERATURE) | |
| response += "\n\nSources:\n" + "\n".join(citations) | |
| return response | |
| except Exception as e: | |
| st.error(f"Error generating answer: {e}") | |
| return "Error generating response." | |
| def get_stats(self) -> Dict: | |
| """to get knowledge base stats""" | |
| if not self.documents: | |
| return {"total": 0} | |
| stats = {"total": len(self.documents)} | |
| doc_types = {} | |
| for doc in self.documents: | |
| doc_type = doc['type'] | |
| doc_types[doc_type] = doc_types.get(doc_type, 0) + 1 | |
| stats.update(doc_types) | |
| return stats |