aniket47/document-intelligence-chatbot
0
1import numpy as np2import pickle3import os4from typing import List, Dict, Tuple5import json6import re7from collections import Counter8import math9import config10 11# Import torch for device handling12try:13 import torch14 TORCH_AVAILABLE = True15except ImportError:16 TORCH_AVAILABLE = False17 18# Import Hugging Face client19try:20 from .huggingface_client import HuggingFaceEmbeddingModel21 HUGGINGFACE_AVAILABLE = True22except ImportError:23 HUGGINGFACE_AVAILABLE = False24 25# Fallback to sentence transformers26try:27 import faiss28 from sentence_transformers import SentenceTransformer29 SENTENCE_TRANSFORMERS_AVAILABLE = True30except ImportError:31 SENTENCE_TRANSFORMERS_AVAILABLE = False32 print("Sentence transformers not available. Using TF-IDF fallback.")33 34class VectorStore:35 """36 Vector store using Sentence Transformers for embeddings and FAISS for similarity search37 """38 39 def __init__(self, model_name: str = None, index_path: str = "vector_index"):40 self.model_name = model_name or config.EMBEDDING_MODEL41 self.index_path = index_path42 self.embedding_model = None43 self.index = None44 self.documents = []45 self.dimension = None46 self.use_huggingface = HUGGINGFACE_AVAILABLE47 self.use_sentence_transformers = SENTENCE_TRANSFORMERS_AVAILABLE48 49 if self.use_huggingface:50 self._load_huggingface_model()51 elif self.use_sentence_transformers:52 self._load_sentence_transformer_model()53 else:54 self._init_simple_search()55 56 def _load_huggingface_model(self):57 """Load the Hugging Face embedding model"""58 try:59 self.embedding_model = HuggingFaceEmbeddingModel(self.model_name)60 # Get dimension61 self.dimension = self.embedding_model.get_dimension()62 print(f"Loaded HuggingFace embedding model: {self.model_name} (dimension: {self.dimension})")63 except Exception as e:64 print(f"Error loading HuggingFace model: {str(e)}")65 self.use_huggingface = False66 if self.use_sentence_transformers:67 self._load_sentence_transformer_model()68 else:69 self._init_simple_search()70 71 def _load_sentence_transformer_model(self):72 """Load the sentence transformer model for embeddings"""73 try:74 # Load with careful device handling - let the library handle device assignment75 self.embedding_model = SentenceTransformer(76 self.model_name, 77 device=None, # Let the library choose the best device78 trust_remote_code=True79 )80 81 # Get dimension from a sample embedding82 sample_embedding = self.embedding_model.encode(["sample"])83 self.dimension = sample_embedding.shape[1] if hasattr(sample_embedding, 'shape') else len(sample_embedding)84 print(f"Loaded sentence transformer model: {self.model_name} (dimension: {self.dimension})")85 except Exception as e:86 print(f"Error loading sentence transformer model: {str(e)}")87 self.use_sentence_transformers = False88 self._init_simple_search()89 90 def _preprocess_text(self, text: str) -> List[str]:91 """Simple text preprocessing for TF-IDF"""92 # Convert to lowercase and remove punctuation93 text = re.sub(r'[^\w\s]', ' ', text.lower())94 # Split into words and remove empty strings95 words = [word for word in text.split() if len(word) > 2]96 return words97 98 def _compute_tf(self, words: List[str]) -> Dict[str, float]:99 """Compute term frequency"""100 word_count = len(words)101 tf_dict = {}102 for word in words:103 tf_dict[word] = tf_dict.get(word, 0) + 1104 # Normalize by total word count105 for word in tf_dict:106 tf_dict[word] = tf_dict[word] / word_count107 return tf_dict108 109 def _compute_idf(self):110 """Compute inverse document frequency for all terms"""111 N = len(self.documents)112 all_words = set()113 for doc in self.documents:114 words = self._preprocess_text(doc['text'])115 all_words.update(set(words))116 117 for word in all_words:118 containing_docs = sum(1 for doc in self.documents 119 if word in self._preprocess_text(doc['text']))120 self.idf_scores[word] = math.log(N / containing_docs) if containing_docs > 0 else 0121 122 def _compute_tfidf_similarity(self, query: str, doc_text: str) -> float:123 """Compute TF-IDF cosine similarity between query and document"""124 query_words = self._preprocess_text(query)125 doc_words = self._preprocess_text(doc_text)126 127 if not query_words or not doc_words:128 return 0.0129 130 query_tf = self._compute_tf(query_words)131 doc_tf = self._compute_tf(doc_words)132 133 # Get all unique words134 all_words = set(query_words + doc_words)135 136 # Compute TF-IDF vectors137 query_vector = []138 doc_vector = []139 140 for word in all_words:141 idf = self.idf_scores.get(word, 0)142 query_tfidf = query_tf.get(word, 0) * idf143 doc_tfidf = doc_tf.get(word, 0) * idf144 query_vector.append(query_tfidf)145 doc_vector.append(doc_tfidf)146 147 # Compute cosine similarity148 if not query_vector or not doc_vector:149 return 0.0150 151 dot_product = sum(a * b for a, b in zip(query_vector, doc_vector))152 query_norm = math.sqrt(sum(a * a for a in query_vector))153 doc_norm = math.sqrt(sum(a * a for a in doc_vector))154 155 if query_norm == 0 or doc_norm == 0:156 return 0.0157 158 return dot_product / (query_norm * doc_norm)159 160 def _init_simple_search(self):161 """Initialize simple TF-IDF search"""162 self.vocabulary = {}163 self.idf_scores = {}164 print("Initialized simple TF-IDF search (advanced embeddings not available)")165 166 def create_embeddings(self, texts: List[str]) -> np.ndarray:167 """Create embeddings for a list of texts"""168 if self.use_huggingface or self.use_sentence_transformers:169 try:170 embeddings = self.embedding_model.encode(texts)171 if hasattr(embeddings, 'numpy'):172 embeddings = embeddings.numpy()173 return embeddings.astype('float32')174 except Exception as e:175 print(f"Error creating embeddings, falling back to simple search: {str(e)}")176 self.use_huggingface = False177 self.use_sentence_transformers = False178 self._init_simple_search()179 180 # Return dummy embeddings for simple search181 return np.zeros((len(texts), 100), dtype='float32')182 183 def initialize_index(self):184 """Initialize FAISS index"""185 if not (self.use_huggingface or self.use_sentence_transformers):186 return187 188 if self.dimension is None:189 raise Exception("Embedding model not properly loaded")190 191 # Use IndexFlatIP for cosine similarity (Inner Product)192 self.index = faiss.IndexFlatIP(self.dimension)193 print(f"Initialized FAISS index with dimension {self.dimension}")194 195 def add_documents(self, chunks: List[Dict]):196 """Add document chunks to the vector store"""197 if not chunks:198 return199 200 # Store documents with metadata201 for i, chunk in enumerate(chunks):202 self.documents.append({203 'id': len(self.documents),204 'text': chunk['text'],205 'metadata': chunk['metadata'],206 'embedding_id': len(self.documents)207 })208 209 if self.use_huggingface or self.use_sentence_transformers:210 # Initialize index if not done211 if self.index is None:212 self.initialize_index()213 214 # Extract texts for embedding215 texts = [chunk['text'] for chunk in chunks]216 217 # Create embeddings218 embeddings = self.create_embeddings(texts)219 220 # Normalize embeddings for cosine similarity221 faiss.normalize_L2(embeddings)222 223 # Add to FAISS index224 self.index.add(embeddings)225 226 print(f"Added {len(chunks)} document chunks to FAISS vector store")227 else:228 # For simple search, compute IDF scores229 self._compute_idf()230 print(f"Added {len(chunks)} document chunks to simple vector store")231 232 def search(self, query: str, k: int = 5, similarity_threshold: float = 0.0) -> List[Dict]:233 """Search for similar documents using semantic similarity with very low threshold"""234 if len(self.documents) == 0:235 return []236 237 if (self.use_huggingface or self.use_sentence_transformers) and self.index is not None:238 return self._advanced_search(query, k, similarity_threshold)239 else:240 return self._simple_search(query, k, similarity_threshold)241 242 def _advanced_search(self, query: str, k: int, similarity_threshold: float) -> List[Dict]:243 """Advanced search using FAISS and sentence transformers"""244 # Create query embedding245 query_embedding = self.create_embeddings([query])246 247 # Normalize for cosine similarity248 faiss.normalize_L2(query_embedding)249 250 # Search in FAISS index251 scores, indices = self.index.search(query_embedding, min(k, len(self.documents)))252 253 results = []254 for i, (score, idx) in enumerate(zip(scores[0], indices[0])):255 # Filter by similarity threshold256 if score >= similarity_threshold and idx < len(self.documents):257 result = {258 'document': self.documents[idx],259 'score': float(score),260 'rank': i + 1261 }262 results.append(result)263 264 return results265 266 def _simple_search(self, query: str, k: int, similarity_threshold: float) -> List[Dict]:267 """Simple search using improved TF-IDF similarity with better matching"""268 if not self.documents:269 return []270 271 # Compute similarities272 similarities = []273 for doc in self.documents:274 # Calculate multiple similarity scores for better matching275 tfidf_similarity = self._compute_tfidf_similarity(query, doc['text'])276 keyword_similarity = self._compute_keyword_similarity(query, doc['text'])277 combined_similarity = max(tfidf_similarity, keyword_similarity * 0.7) # Boost keyword matches278 279 similarities.append({280 'document': doc,281 'score': combined_similarity,282 'rank': 0 # Will be set after sorting283 })284 285 # Sort by similarity score286 similarities.sort(key=lambda x: x['score'], reverse=True)287 288 # Always return results, ignore similarity threshold for TF-IDF fallback289 results = []290 for i, result in enumerate(similarities[:k]):291 result['rank'] = i + 1292 results.append(result)293 294 return results295 296 def _compute_keyword_similarity(self, query: str, text: str) -> float:297 """Compute simple keyword-based similarity"""298 query_words = set(query.lower().split())299 text_words = set(text.lower().split())300 301 if not query_words:302 return 0.0303 304 # Calculate Jaccard similarity305 intersection = query_words.intersection(text_words)306 union = query_words.union(text_words)307 308 if not union:309 return 0.0310 311 return len(intersection) / len(union)312 313 def save_index(self):314 """Save vector store to disk"""315 try:316 if (self.use_huggingface or self.use_sentence_transformers) and self.index is not None:317 # Save FAISS index318 faiss.write_index(self.index, f"{self.index_path}.faiss")319 320 # Save documents and metadata321 with open(f"{self.index_path}_docs.pkl", "wb") as f:322 pickle.dump({323 'documents': self.documents,324 'dimension': self.dimension,325 'model_name': self.model_name,326 'use_huggingface': self.use_huggingface,327 'use_sentence_transformers': self.use_sentence_transformers,328 'vocabulary': getattr(self, 'vocabulary', {}),329 'idf_scores': getattr(self, 'idf_scores', {})330 }, f)331 332 print(f"Saved vector index to {self.index_path}")333 except Exception as e:334 print(f"Error saving index: {str(e)}")335 336 def load_index(self):337 """Load vector store from disk"""338 try:339 if os.path.exists(f"{self.index_path}_docs.pkl"):340 # Load documents and metadata341 with open(f"{self.index_path}_docs.pkl", "rb") as f:342 data = pickle.load(f)343 self.documents = data['documents']344 self.dimension = data.get('dimension')345 self.vocabulary = data.get('vocabulary', {})346 self.idf_scores = data.get('idf_scores', {})347 stored_use_hf = data.get('use_huggingface', False)348 stored_use_st = data.get('use_sentence_transformers', data.get('use_advanced', True))349 350 # Load FAISS index if available and we're using embeddings351 if ((self.use_huggingface or self.use_sentence_transformers) and 352 (stored_use_hf or stored_use_st) and 353 os.path.exists(f"{self.index_path}.faiss")):354 self.index = faiss.read_index(f"{self.index_path}.faiss")355 356 print(f"Loaded vector index from {self.index_path}")357 return True358 except Exception as e:359 print(f"Error loading index: {str(e)}")360 361 return False362 363 def clear_index(self):364 """Clear the current index and documents"""365 self.index = None366 self.documents = []367 self.vocabulary = {}368 self.idf_scores = {}369 print("Cleared vector index")370 371 def get_stats(self) -> Dict:372 """Get statistics about the vector store"""373 return {374 'total_documents': len(self.documents),375 'index_size': self.index.ntotal if ((self.use_huggingface or self.use_sentence_transformers) and self.index) else len(self.documents),376 'dimension': self.dimension,377 'model_name': self.model_name,378 'search_type': 'HuggingFace Embeddings + FAISS' if self.use_huggingface else 'Sentence Transformers + FAISS' if self.use_sentence_transformers else 'Simple TF-IDF'379 }