Spaces:
Sleeping
Sleeping
| """ | |
| RAG (Retrieval Augmented Generation) utilities using FAISS and embeddings | |
| """ | |
| import numpy as np | |
| from typing import List, Dict, Any | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| class SimpleRAGStore: | |
| """ | |
| Simple RAG implementation using FAISS for vector similarity search | |
| """ | |
| def __init__(self): | |
| """Initialize the RAG store""" | |
| self.documents: List[str] = [] | |
| self.embeddings: List[np.ndarray] = [] | |
| self.index = None | |
| self._model = None | |
| def _get_model(self): | |
| """Lazy load the sentence transformer model""" | |
| if self._model is None: | |
| try: | |
| from sentence_transformers import SentenceTransformer | |
| self._model = SentenceTransformer('all-MiniLM-L6-v2') | |
| logger.info("Loaded sentence transformer model") | |
| except Exception as e: | |
| logger.error(f"Failed to load sentence transformer: {e}") | |
| raise | |
| return self._model | |
| def add_documents(self, documents: List[str]) -> None: | |
| """ | |
| Add documents to the RAG store and build FAISS index. | |
| Args: | |
| documents: List of document strings to add | |
| """ | |
| import faiss | |
| if not documents: | |
| logger.warning("No documents provided to add") | |
| return | |
| self.documents.extend(documents) | |
| # Generate embeddings | |
| model = self._get_model() | |
| new_embeddings = model.encode(documents, show_progress_bar=False) | |
| self.embeddings.extend(new_embeddings) | |
| # Build or update FAISS index | |
| embeddings_array = np.array(self.embeddings).astype('float32') | |
| dimension = embeddings_array.shape[1] | |
| if self.index is None: | |
| self.index = faiss.IndexFlatL2(dimension) | |
| self.index.add(embeddings_array) | |
| logger.info(f"Added {len(documents)} documents to RAG store") | |
| def search(self, query: str, top_k: int = 3) -> List[Dict[str, Any]]: | |
| """ | |
| Search for similar documents using the query. | |
| Args: | |
| query: Search query string | |
| top_k: Number of top results to return | |
| Returns: | |
| List of search results with scores | |
| """ | |
| if self.index is None or len(self.documents) == 0: | |
| logger.warning("No documents in RAG store") | |
| return [] | |
| # Encode query | |
| model = self._get_model() | |
| query_embedding = model.encode([query], show_progress_bar=False) | |
| query_embedding = np.array(query_embedding).astype('float32') | |
| # Search FAISS index | |
| top_k = min(top_k, len(self.documents)) | |
| distances, indices = self.index.search(query_embedding, top_k) | |
| # Format results | |
| results = [] | |
| for i, (distance, idx) in enumerate(zip(distances[0], indices[0])): | |
| if idx < len(self.documents): | |
| # Convert L2 distance to similarity score (inverse relationship) | |
| similarity_score = 1.0 / (1.0 + float(distance)) | |
| results.append({ | |
| "rank": i + 1, | |
| "document": self.documents[idx], | |
| "score": round(similarity_score, 4), | |
| "distance": float(distance) | |
| }) | |
| return results | |
| def clear(self) -> None: | |
| """Clear all documents and reset the index""" | |
| self.documents = [] | |
| self.embeddings = [] | |
| self.index = None | |
| logger.info("Cleared RAG store") | |
| def create_rag_store(documents: List[str]) -> SimpleRAGStore: | |
| """ | |
| Factory function to create and populate a RAG store. | |
| Args: | |
| documents: List of documents to add to store | |
| Returns: | |
| Initialized SimpleRAGStore instance | |
| """ | |
| store = SimpleRAGStore() | |
| if documents: | |
| store.add_documents(documents) | |
| return store | |
| def semantic_search(query: str, documents: List[str], top_k: int = 3) -> List[Dict[str, Any]]: | |
| """ | |
| Perform semantic search on a list of documents. | |
| Args: | |
| query: Search query | |
| documents: List of documents to search | |
| top_k: Number of results to return | |
| Returns: | |
| List of search results | |
| """ | |
| store = create_rag_store(documents) | |
| return store.search(query, top_k) | |