#!/usr/bin/env python3 """ # Simplified RAG System for Hugging Face Spaces This module provides a comprehensive Retrieval-Augmented Generation (RAG) system using: - **FAISS** for efficient vector storage and similarity search - **BM25** for sparse retrieval and keyword matching - **Hybrid Search** combining both dense and sparse methods - **Qwen 2.5 1.5B** for intelligent response generation - **Thread Safety** for concurrent document loading ## Architecture Overview The RAG system follows a modular design with these key components: 1. **Document Processing**: PDF extraction and intelligent chunking 2. **Vector Storage**: FAISS index for high-dimensional embeddings 3. **Sparse Retrieval**: BM25 for keyword-based search 4. **Hybrid Search**: Combines dense and sparse methods for optimal results 5. **Response Generation**: LLM-based answer synthesis with context 6. **Thread Safety**: Concurrent document loading with proper locking ## Key Features - 🔍 **Multi-Method Search**: Hybrid, dense, and sparse retrieval options - 📊 **Performance Metrics**: Confidence scores and response times - 🔒 **Thread Safety**: Safe concurrent document loading - 💾 **Persistence**: Automatic index saving and loading - 🎯 **Smart Fallbacks**: Graceful model loading with alternatives - 📈 **Scalable**: Efficient handling of large document collections ## Usage Example ```python # Initialize the RAG system rag = SimpleRAGSystem() # Add documents rag.add_document("document.pdf", "Document Name") # Query the system response = rag.query("What is the main topic?", method="hybrid", top_k=5) print(response.answer) ``` """ import os import pickle import json import time from typing import List, Dict, Optional, Tuple from dataclasses import dataclass import numpy as np import torch from loguru import logger import threading # Import required libraries for AI/ML functionality from sentence_transformers import SentenceTransformer from rank_bm25 import BM25Okapi import faiss from transformers import AutoTokenizer, AutoModelForCausalLM # Import guard rail system from guard_rails import GuardRailSystem, GuardRailConfig, GuardRailResult # Import HF Spaces configuration try: from hf_spaces_config import get_hf_config, is_hf_spaces HF_SPACES_AVAILABLE = True except ImportError: HF_SPACES_AVAILABLE = False logger.warning("HF Spaces configuration not available") # ============================================================================= # DATA STRUCTURES # ============================================================================= @dataclass class DocumentChunk: """ Represents a document chunk with metadata Attributes: text: The actual text content of the chunk doc_id: Unique identifier for the source document filename: Name of the source file chunk_id: Unique identifier for this specific chunk chunk_size: Target size used for chunking """ text: str doc_id: str filename: str chunk_id: str chunk_size: int @dataclass class SearchResult: """ Represents a search result with scoring information Attributes: text: The retrieved text content score: Combined relevance score doc_id: Source document identifier filename: Source file name search_method: Method used for retrieval (dense/sparse/hybrid) dense_score: Vector similarity score (if applicable) sparse_score: Keyword matching score (if applicable) """ text: str score: float doc_id: str filename: str search_method: str dense_score: Optional[float] = None sparse_score: Optional[float] = None @dataclass class RAGResponse: """ Represents a complete RAG system response Attributes: answer: Generated answer text confidence: Confidence score for the response search_results: List of retrieved documents method_used: Search method that was used response_time: Time taken to generate response query: Original user query """ answer: str confidence: float search_results: List[SearchResult] method_used: str response_time: float query: str # ============================================================================= # MAIN RAG SYSTEM CLASS # ============================================================================= class SimpleRAGSystem: """ Simplified RAG system for Hugging Face Spaces This class provides a complete RAG implementation with: - Document ingestion and processing - Vector and sparse search capabilities - Response generation using language models - Thread-safe concurrent operations - Persistent storage and retrieval """ def __init__( self, embedding_model: str = "all-MiniLM-L6-v2", generative_model: str = "Qwen/Qwen2.5-1.5B-Instruct", chunk_sizes: List[int] = None, vector_store_path: str = "./vector_store", enable_guard_rails: bool = True, guard_rail_config: GuardRailConfig = None, ): """ Initialize the RAG system with specified models and configuration Args: embedding_model: Sentence transformer model for embeddings generative_model: Language model for response generation chunk_sizes: List of chunk sizes for document processing vector_store_path: Path for storing FAISS index and metadata enable_guard_rails: Whether to enable guard rail system guard_rail_config: Configuration for guard rail system """ self.embedding_model = embedding_model self.generative_model = generative_model self.chunk_sizes = chunk_sizes or [100, 400] # Default chunk sizes self.vector_store_path = vector_store_path self.enable_guard_rails = enable_guard_rails # Initialize core components self.embedder = None # Sentence transformer for embeddings self.tokenizer = None # Tokenizer for language model self.model = None # Language model for generation self.faiss_index = None # FAISS index for vector search self.bm25 = None # BM25 for sparse search self.documents = [] # List of processed documents self.chunks = [] # List of document chunks self._lock = threading.Lock() # Thread safety for concurrent loading # Initialize guard rail system if self.enable_guard_rails: self.guard_rails = GuardRailSystem(guard_rail_config) logger.info("Guard rail system enabled") else: self.guard_rails = None logger.info("Guard rail system disabled") # Create vector store directory for persistence os.makedirs(vector_store_path, exist_ok=True) # Set up HF Spaces configuration if available if HF_SPACES_AVAILABLE: try: hf_config = get_hf_config() if is_hf_spaces(): logger.info( "🌐 HF Spaces environment detected - using optimized configuration" ) # Cache directories are automatically set up by hf_config else: logger.info("💻 Local development environment detected") except Exception as e: logger.warning(f"HF Spaces configuration failed: {e}") # Load or initialize system components self._load_models() self._load_or_create_index() logger.info("Simple RAG system initialized successfully!") def _load_models(self): """ Load embedding and generative models with fallback handling This method: 1. Loads the sentence transformer for embeddings 2. Attempts to load the primary language model (Qwen) 3. Falls back to distilgpt2 if primary model fails 4. Configures tokenizers and model settings """ try: # Get cache directory from HF Spaces config if available cache_dir = None if HF_SPACES_AVAILABLE: try: hf_config = get_hf_config() cache_dir = hf_config.cache_dirs.get("transformers_cache") if cache_dir: logger.info(f"Using HF Spaces cache directory: {cache_dir}") except Exception as e: logger.warning(f"Could not get HF Spaces cache directory: {e}") # Load embedding model for document vectorization if cache_dir: self.embedder = SentenceTransformer( self.embedding_model, cache_folder=cache_dir ) else: self.embedder = SentenceTransformer(self.embedding_model) self.vector_size = self.embedder.get_sentence_embedding_dimension() # Load generative model with fallback strategy model_loaded = False # Try loading Qwen model first (primary choice) try: self.tokenizer = AutoTokenizer.from_pretrained( self.generative_model, trust_remote_code=True, padding_side="left", # Important for generation cache_dir=cache_dir, ) # Load model with explicit CPU configuration for deployment compatibility self.model = AutoModelForCausalLM.from_pretrained( self.generative_model, trust_remote_code=True, torch_dtype=torch.float32, # Use float32 for CPU compatibility device_map=None, # Let PyTorch handle device placement low_cpu_mem_usage=False, # Disable for better compatibility cache_dir=cache_dir, ) # Move to CPU explicitly for deployment environments self.model = self.model.to("cpu") model_loaded = True except Exception as e: logger.warning(f"Failed to load Qwen model: {e}") # Fallback to distilgpt2 if Qwen fails if not model_loaded: logger.info("Falling back to distilgpt2...") self.generative_model = "distilgpt2" try: self.tokenizer = AutoTokenizer.from_pretrained( self.generative_model, trust_remote_code=True, padding_side="left", ) self.model = AutoModelForCausalLM.from_pretrained( self.generative_model, trust_remote_code=True, ) # Ensure fallback model is also on CPU self.model = self.model.to("cpu") model_loaded = True except Exception as e: logger.error(f"Failed to load distilgpt2: {e}") raise Exception("Could not load any generative model") # Configure tokenizer settings for generation if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token self.tokenizer.pad_token_id = self.tokenizer.eos_token_id logger.info(f"✅ Models loaded successfully") logger.info(f" - Embedding: {self.embedding_model}") logger.info(f" - Generative: {self.generative_model}") except Exception as e: logger.error(f"❌ Failed to load models: {e}") raise def _load_or_create_index(self): """ Load existing FAISS index or create a new one This method: 1. Checks for existing index files 2. Loads existing index and metadata if available 3. Creates new index if none exists 4. Rebuilds BM25 index from loaded chunks """ faiss_path = os.path.join(self.vector_store_path, "faiss_index.bin") metadata_path = os.path.join(self.vector_store_path, "metadata.pkl") if os.path.exists(faiss_path) and os.path.exists(metadata_path): # Load existing index and metadata try: self.faiss_index = faiss.read_index(faiss_path) with open(metadata_path, "rb") as f: metadata = pickle.load(f) self.documents = metadata.get("documents", []) self.chunks = metadata.get("chunks", []) # Rebuild BM25 index from loaded chunks if self.chunks: texts = [chunk.text for chunk in self.chunks] tokenized_texts = [text.lower().split() for text in texts] self.bm25 = BM25Okapi(tokenized_texts) logger.info(f"✅ Loaded existing index with {len(self.chunks)} chunks") except Exception as e: logger.warning(f"Failed to load existing index: {e}") self._create_new_index() else: self._create_new_index() def _create_new_index(self): """Create new FAISS index with appropriate configuration""" vector_size = self.embedder.get_sentence_embedding_dimension() # Use Inner Product for cosine similarity (normalized vectors) self.faiss_index = faiss.IndexFlatIP(vector_size) self.bm25 = None logger.info(f"✅ Created new FAISS index with dimension {vector_size}") def _save_index(self): """ Save FAISS index and metadata for persistence This ensures that the system state is preserved across restarts. """ try: # Save FAISS index faiss_path = os.path.join(self.vector_store_path, "faiss_index.bin") faiss.write_index(self.faiss_index, faiss_path) # Save metadata including documents and chunks metadata_path = os.path.join(self.vector_store_path, "metadata.pkl") metadata = {"documents": self.documents, "chunks": self.chunks} with open(metadata_path, "wb") as f: pickle.dump(metadata, f) logger.info("✅ Index saved successfully") except Exception as e: logger.error(f"❌ Failed to save index: {e}") def add_document(self, file_path: str, filename: str) -> bool: """ Add a document to the RAG system with thread safety This method: 1. Processes the PDF document into chunks 2. Adds document metadata to the system 3. Updates embeddings and BM25 index 4. Saves the updated index Args: file_path: Path to the PDF file filename: Name of the file for reference Returns: True if successful, False otherwise """ try: from pdf_processor import SimplePDFProcessor # Process the document using the PDF processor processor = SimplePDFProcessor() processed_doc = processor.process_document(file_path, self.chunk_sizes) # Thread-safe document addition using lock with self._lock: # Add document metadata to the system self.documents.append( { "filename": filename, "title": processed_doc.title, "author": processed_doc.author, "file_path": file_path, } ) # Add all chunks from the processed document for chunk in processed_doc.chunks: self.chunks.append(chunk) # Update search indices with new content self._update_embeddings() self._update_bm25() # Persist the updated index self._save_index() logger.info( f"✅ Added document: {filename} ({len(processed_doc.chunks)} chunks)" ) return True except Exception as e: logger.error(f"❌ Failed to add document {filename}: {e}") return False def _update_embeddings(self): """ Update FAISS index with new embeddings This method: 1. Extracts text from all chunks 2. Generates embeddings using the sentence transformer 3. Adds embeddings to the FAISS index """ if not self.chunks: return # Generate embeddings for all chunks texts = [chunk.text for chunk in self.chunks] embeddings = self.embedder.encode(texts, show_progress_bar=False) # Add embeddings to FAISS index self.faiss_index.add(embeddings.astype("float32")) def _update_bm25(self): """ Update BM25 index with new chunks This method rebuilds the BM25 index with all current chunks for keyword-based search functionality. """ if not self.chunks: return # Rebuild BM25 with all chunks texts = [chunk.text for chunk in self.chunks] tokenized_texts = [text.lower().split() for text in texts] self.bm25 = BM25Okapi(tokenized_texts) def search( self, query: str, method: str = "hybrid", top_k: int = 5 ) -> List[SearchResult]: """ Search for relevant documents using specified method This method supports three search strategies: - **dense**: Vector similarity search using FAISS - **sparse**: Keyword matching using BM25 - **hybrid**: Combines both methods for optimal results Args: query: Search query string method: Search method (hybrid, dense, sparse) top_k: Number of results to return Returns: List of search results with scores and metadata """ if not self.chunks: return [] results = [] # Perform dense search (vector similarity) if method == "dense" or method == "hybrid": # Generate query embedding query_embedding = self.embedder.encode([query]) # Search FAISS index scores, indices = self.faiss_index.search( query_embedding.astype("float32"), min(top_k, len(self.chunks)) ) # Process dense search results for score, idx in zip(scores[0], indices[0]): if idx < len(self.chunks): chunk = self.chunks[idx] results.append( SearchResult( text=chunk.text, score=float(score), doc_id=chunk.doc_id, filename=chunk.filename, search_method="dense", dense_score=float(score), ) ) # Perform sparse search (keyword matching) if method == "sparse" or method == "hybrid": if self.bm25: # Tokenize query for BM25 tokenized_query = query.lower().split() bm25_scores = self.bm25.get_scores(tokenized_query) # Get top BM25 results top_indices = np.argsort(bm25_scores)[::-1][:top_k] # Process sparse search results for idx in top_indices: if idx < len(self.chunks): chunk = self.chunks[idx] score = float(bm25_scores[idx]) # Check if result already exists (for hybrid search) existing_result = next( ( r for r in results if r.doc_id == chunk.doc_id and r.text == chunk.text ), None, ) if existing_result: # Update existing result with sparse score existing_result.sparse_score = score if method == "hybrid": # Combine scores for hybrid search existing_result.score = ( existing_result.dense_score + score ) / 2 else: # Add new sparse result results.append( SearchResult( text=chunk.text, score=score, doc_id=chunk.doc_id, filename=chunk.filename, search_method="sparse", sparse_score=score, ) ) # Sort by score and return top_k results results.sort(key=lambda x: x.score, reverse=True) return results[:top_k] def generate_response(self, query: str, context: str) -> str: """ Generate response using the language model This method: 1. Prepares a prompt with context and query 2. Uses the appropriate chat template for the model 3. Generates a response with controlled parameters 4. Handles model-specific response formatting Args: query: User's question context: Retrieved context from search Returns: Generated response text """ try: # Prepare prompt based on model capabilities if hasattr(self.tokenizer, "apply_chat_template"): # Use chat template for modern models like Qwen messages = [ { "role": "system", "content": "You are a helpful AI assistant. Use the provided context to answer the user's question accurately and concisely. If the context doesn't contain enough information to answer the question, say so.", }, { "role": "user", "content": f"Context: {context}\n\nQuestion: {query}", }, ] prompt = self.tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) else: # Fallback for non-chat models prompt = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:" # Tokenize input with appropriate settings tokenized = self.tokenizer( prompt, return_tensors="pt", truncation=True, max_length=1024, # Limit input length padding=True, return_attention_mask=True, ) # Generate response with controlled parameters with torch.no_grad(): try: outputs = self.model.generate( tokenized.input_ids, attention_mask=tokenized.attention_mask, max_new_tokens=512, # Limit response length num_return_sequences=1, temperature=0.7, # Balance creativity and coherence do_sample=True, # Enable sampling for more natural responses pad_token_id=self.tokenizer.pad_token_id, eos_token_id=self.tokenizer.eos_token_id, ) except RuntimeError as e: if "Half" in str(e): # Handle half-precision compatibility issues logger.warning( "Half precision not supported on CPU, converting to float32" ) # Convert model to float32 self.model = self.model.float() outputs = self.model.generate( tokenized.input_ids, attention_mask=tokenized.attention_mask, max_new_tokens=512, num_return_sequences=1, temperature=0.7, do_sample=True, pad_token_id=self.tokenizer.pad_token_id, eos_token_id=self.tokenizer.eos_token_id, ) else: raise e # Decode the generated response response = self.tokenizer.decode(outputs[0], skip_special_tokens=True) # Extract only the generated part (remove input prompt) if hasattr(self.tokenizer, "apply_chat_template"): # Handle chat model response formatting if "<|im_start|>assistant" in response: response = response.split("<|im_start|>assistant")[-1] if "<|im_end|>" in response: response = response.split("<|im_end|>")[0] else: # Handle standard model response formatting response = response[len(prompt) :] return response.strip() except Exception as e: logger.error(f"Error generating response: {e}") return f"Error generating response: {str(e)}" def query( self, query: str, method: str = "hybrid", top_k: int = 5, user_id: str = "anonymous", ) -> RAGResponse: """ Complete RAG query pipeline with guard rail protection This method orchestrates the entire RAG process with safety checks: 1. Validates input using guard rails 2. Searches for relevant documents 3. Combines context from search results 4. Generates a response using the language model 5. Validates output using guard rails 6. Calculates confidence and timing metrics Args: query: User's question method: Search method to use top_k: Number of search results to use user_id: User identifier for rate limiting and tracking Returns: Complete RAG response with answer, metadata, and metrics """ start_time = time.time() # ============================================================================= # INPUT VALIDATION WITH GUARD RAILS # ============================================================================= if self.enable_guard_rails and self.guard_rails: # Validate input using guard rails input_validation = self.guard_rails.validate_input(query, user_id) if not input_validation.passed: logger.warning(f"Input validation failed: {input_validation.reason}") if input_validation.blocked: return RAGResponse( answer=f"I cannot process this request: {input_validation.reason}", confidence=0.0, search_results=[], method_used=method, response_time=time.time() - start_time, query=query, ) else: # Warning but continue processing logger.warning( f"Input validation warning: {input_validation.reason}" ) # Sanitize input query = self.guard_rails.sanitize_input(query) # Search for relevant documents search_results = self.search(query, method, top_k) # Handle case where no relevant documents found if not search_results: return RAGResponse( answer="I couldn't find any relevant information to answer your question.", confidence=0.0, search_results=[], method_used=method, response_time=time.time() - start_time, query=query, ) # Combine context from search results context = "\n\n".join([result.text for result in search_results]) # Generate response using the language model answer = self.generate_response(query, context) # Calculate confidence based on search result scores confidence = np.mean([result.score for result in search_results]) # ============================================================================= # OUTPUT VALIDATION WITH GUARD RAILS # ============================================================================= if self.enable_guard_rails and self.guard_rails: # Validate output using guard rails output_validation = self.guard_rails.validate_output( answer, confidence, context ) if not output_validation.passed: logger.warning(f"Output validation failed: {output_validation.reason}") if output_validation.blocked: return RAGResponse( answer="I cannot provide this response due to safety concerns.", confidence=0.0, search_results=search_results, method_used=method, response_time=time.time() - start_time, query=query, ) else: # Warning but continue with response logger.warning( f"Output validation warning: {output_validation.reason}" ) # Sanitize output answer = self.guard_rails.sanitize_output(answer) # Create and return complete response return RAGResponse( answer=answer, confidence=confidence, search_results=search_results, method_used=method, response_time=time.time() - start_time, query=query, ) def get_stats(self) -> Dict: """ Get system statistics and configuration information Returns: Dictionary containing system metrics and settings """ return { "total_documents": len(self.documents), "total_chunks": len(self.chunks), "vector_size": ( self.embedder.get_sentence_embedding_dimension() if self.embedder else 0 ), "model_name": self.generative_model, "embedding_model": self.embedding_model, "chunk_sizes": self.chunk_sizes, } def clear(self): """ Clear all documents and reset the system This method: 1. Clears all documents and chunks 2. Creates a new FAISS index 3. Saves the empty state """ self.documents = [] self.chunks = [] self._create_new_index() self._save_index() logger.info("✅ System cleared successfully")