aniket47/document-intelligence-chatbot
0
1"""2Local Hugging Face model integration with automatic model downloading3"""4 5import os6import torch7from typing import List, Dict, Optional8import config9import warnings10 11# Suppress some warnings for cleaner output12warnings.filterwarnings("ignore", category=UserWarning, module="transformers")13 14class HuggingFaceClient:15 """16 Client for local Hugging Face models with automatic downloading17 """18 19 def __init__(self, model_name: str = None, cache_dir: str = None):20 self.model_name = model_name or config.CHAT_MODEL21 self.cache_dir = cache_dir or config.MODEL_CACHE_DIR22 self.max_length = config.MODEL_MAX_LENGTH23 self.temperature = config.TEMPERATURE24 25 # Create cache directory if it doesn't exist26 os.makedirs(self.cache_dir, exist_ok=True)27 28 # Initialize device29 self.device = self._setup_device()30 31 # Initialize models (will be loaded on first use)32 self.tokenizer = None33 self.model = None34 self.model_type = None # Will be set during loading35 self.is_loaded = False36 37 print(f"HuggingFace Client initialized")38 print(f"Model: {self.model_name}")39 print(f"Cache: {self.cache_dir}")40 print(f"Device: {self.device}")41 42 def _setup_device(self):43 """Setup computation device (CPU/GPU)"""44 if config.DEVICE == "auto":45 if config.USE_CUDA and torch.cuda.is_available():46 device = "cuda"47 print(f"Using GPU: {torch.cuda.get_device_name()}")48 else:49 device = "cpu"50 print("Using CPU")51 else:52 device = config.DEVICE53 54 return device55 56 def _load_model(self):57 """Load the model and tokenizer (downloads automatically if not cached)"""58 if self.is_loaded:59 return True60 61 try:62 print(f"Loading model: {self.model_name}")63 print("This might take a few minutes on first run (downloading model)...")64 65 # Import here to avoid slow startup if not needed66 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM67 68 # Load tokenizer69 self.tokenizer = AutoTokenizer.from_pretrained(70 self.model_name,71 cache_dir=self.cache_dir72 )73 74 # Determine model type and load accordingly75 is_t5_model = "t5" in self.model_name.lower() or "flan" in self.model_name.lower()76 77 if is_t5_model:78 print("Loading T5/FLAN model for text-to-text generation...")79 self.model = AutoModelForSeq2SeqLM.from_pretrained(80 self.model_name,81 cache_dir=self.cache_dir,82 torch_dtype=torch.float32, # T5 works better with float3283 low_cpu_mem_usage=True,84 trust_remote_code=True85 )86 self.model_type = "seq2seq"87 print("T5/FLAN model loaded successfully!")88 else:89 print("Loading causal language model...")90 self.model = AutoModelForCausalLM.from_pretrained(91 self.model_name,92 cache_dir=self.cache_dir,93 torch_dtype=torch.float32,94 low_cpu_mem_usage=True,95 trust_remote_code=True96 )97 self.model_type = "causal"98 99 # Add pad token for causal models100 if self.tokenizer.pad_token is None:101 self.tokenizer.pad_token = self.tokenizer.eos_token102 print("Causal model loaded successfully!")103 104 self.model.eval() # Set to evaluation mode105 self.is_loaded = True106 107 print(f"Model size: ~{self._get_model_size_mb():.1f} MB")108 return True109 110 except Exception as e:111 print(f"Error loading model: {str(e)}")112 print("Model will run in offline mode - document search will still work!")113 self.is_loaded = False114 return False115 116 def _get_model_size_mb(self):117 """Estimate model size in MB"""118 if self.model is None:119 return 0120 121 param_size = 0122 for param in self.model.parameters():123 param_size += param.nelement() * param.element_size()124 125 return param_size / 1024 / 1024126 127 def generate_response(self, query: str, context: str = "", system_prompt: str = "") -> str:128 """Generate a response given a query and context with offline fallback"""129 # Load model on first use130 if not self.is_loaded:131 success = self._load_model()132 if not success:133 # Return offline fallback response134 return self._generate_offline_response(query, context)135 136 try:137 # Prepare the input text based on model type138 if hasattr(self, 'model_type') and self.model_type == "seq2seq":139 # T5/FLAN models work better with instruction-style prompts140 if context:141 # For document-based questions142 context_truncated = context[:800] if len(context) > 800 else context143 144 if any(word in query.lower() for word in ['summarize', 'summary', 'main points', 'key points', 'overview']):145 input_text = f"Summarize the following text: {context_truncated}"146 else:147 input_text = f"Answer the question based on the context.\nContext: {context_truncated}\nQuestion: {query}\nAnswer:"148 else:149 input_text = f"Answer this question: {query}"150 151 # Tokenize for T5152 input_ids = self.tokenizer.encode(input_text, return_tensors="pt", truncation=True, max_length=512)153 154 # Ensure input_ids are on the same device as the model155 if hasattr(self.model, 'device'):156 model_device = next(self.model.parameters()).device157 input_ids = input_ids.to(model_device)158 else:159 input_ids = input_ids.to(self.device)160 161 # Generate with T5/FLAN162 with torch.no_grad():163 outputs = self.model.generate(164 input_ids,165 max_length=200, # Good length for summaries166 min_length=20, # Ensure substantial response167 temperature=0.7,168 do_sample=True,169 pad_token_id=self.tokenizer.pad_token_id,170 eos_token_id=self.tokenizer.eos_token_id,171 num_return_sequences=1,172 no_repeat_ngram_size=3,173 length_penalty=1.0174 )175 176 # Decode T5 response (T5 outputs only the generated text)177 response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)178 179 else:180 # Original logic for causal models (DialoGPT, etc.)181 if context:182 context_truncated = context[:500] if len(context) > 500 else context183 184 if any(word in query.lower() for word in ['summarize', 'summary', 'main points', 'key points', 'overview']):185 input_text = f"Summarize this: {context_truncated}\nSummary:"186 else:187 input_text = f"Context: {context_truncated}\nQuestion: {query}\nAnswer:"188 else:189 input_text = f"Question: {query}\nAnswer:"190 191 # Tokenize input with simpler approach192 input_ids = self.tokenizer.encode(input_text, return_tensors="pt", truncation=True, max_length=300)193 194 # Ensure input_ids are on the same device as the model195 if hasattr(self.model, 'device'):196 model_device = next(self.model.parameters()).device197 input_ids = input_ids.to(model_device)198 else:199 input_ids = input_ids.to(self.device)200 201 # Generate response with causal model202 with torch.no_grad():203 outputs = self.model.generate(204 input_ids,205 max_length=input_ids.shape[1] + 100,206 min_length=input_ids.shape[1] + 5,207 temperature=0.8,208 do_sample=True,209 pad_token_id=self.tokenizer.eos_token_id,210 eos_token_id=self.tokenizer.eos_token_id,211 num_return_sequences=1,212 no_repeat_ngram_size=2,213 repetition_penalty=1.1,214 length_penalty=1.0215 )216 217 # Decode causal model response218 response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)219 220 # Extract only the new generated text for causal models221 if response.startswith(input_text):222 response = response[len(input_text):].strip()223 else:224 # Fallback: try to find the answer part225 for separator in ["Answer:", "Summary:", "\nBot:", "\n"]:226 if separator in response:227 parts = response.split(separator)228 if len(parts) > 1:229 response = parts[-1].strip()230 break231 232 print(f"Extracted response: '{response[:100]}...'")233 234 # Clean up the response235 cleaned_response = self._clean_response(response)236 237 # Debug logging238 print(f"Raw AI response length: {len(response)}")239 print(f"Cleaned AI response length: {len(cleaned_response)}")240 print(f"Cleaned response: '{cleaned_response[:100]}...'")241 242 # Be more lenient - if we have any response, use it243 if cleaned_response and len(cleaned_response.strip()) > 0:244 return cleaned_response245 elif response and len(response.strip()) > 0:246 # Use raw response if cleaning removed too much247 return response.strip()248 else:249 # Try a simple fallback generation250 print("Attempting fallback generation with simpler prompt...")251 return self._try_simple_generation(query, context)252 253 except Exception as e:254 print(f"Error generating response: {str(e)}")255 # Fall back to offline response256 return self._generate_offline_response(query, context)257 258 def _try_simple_generation(self, query: str, context: str = "") -> str:259 """Try a very simple generation as last resort"""260 try:261 # Ultra-simple prompt262 simple_prompt = f"{query}"263 input_ids = self.tokenizer.encode(simple_prompt, return_tensors="pt", max_length=50)264 265 # Ensure input_ids are on the same device as the model266 if hasattr(self.model, 'device'):267 model_device = next(self.model.parameters()).device268 input_ids = input_ids.to(model_device)269 else:270 input_ids = input_ids.to(self.device)271 272 with torch.no_grad():273 outputs = self.model.generate(274 input_ids,275 max_length=input_ids.shape[1] + 30,276 temperature=0.9,277 do_sample=True,278 pad_token_id=self.tokenizer.eos_token_id,279 num_return_sequences=1280 )281 282 response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)283 response = response[len(simple_prompt):].strip()284 285 if response and len(response) > 2:286 return f"AI Response: {response}"287 288 except Exception as e:289 print(f"Simple generation also failed: {e}")290 291 return self._generate_offline_response(query, context)292 293 def _generate_offline_response(self, query: str, context: str = "") -> str:294 """Generate a structured response when AI model is unavailable or gives poor response"""295 # Check if this is being called because model is unavailable or just poor response296 model_available = self.is_loaded297 note_suffix = "*Note: AI model generated poor response - showing raw content*" if model_available else "*Note: AI model unavailable - showing raw content*"298 299 if context:300 if "Relevant information from your documents:" in context:301 # Extract and format document content302 lines = context.split('\n')303 document_info = []304 current_info = ""305 306 for line in lines:307 line = line.strip()308 if line.startswith("From ") and "relevance:" in line:309 if current_info:310 document_info.append(current_info)311 # Extract filename312 filename = line.split("(relevance:")[0].replace("From ", "").strip()313 current_info = f"**From {filename}:**"314 elif line and not line.startswith("Relevant information") and len(line) > 10:315 current_info += f"\n{line}"316 317 if current_info:318 document_info.append(current_info)319 320 if document_info:321 response = "Based on your uploaded documents:\n\n"322 for info in document_info[:2]: # Show top 2 sources323 response += f"{info}\n\n"324 response += f"\n{note_suffix}"325 return response326 327 elif "Web search results:" in context:328 # Format web search results329 lines = context.split('\n')330 search_results = []331 332 for line in lines:333 if line.strip() and not line.startswith('Web search results:'):334 search_results.append(line.strip())335 336 if search_results:337 response = "Based on web search results:\n\n"338 for i, result in enumerate(search_results[:3], 1):339 response += f"{i}. {result}\n"340 response += f"\n{note_suffix}"341 return response342 343 # No context or fallback case344 if model_available:345 return (f"I received your question: '{query}'\n\n"346 f"I'm having trouble generating a good response right now. "347 f"This might be due to the complexity of the question or model limitations.\n\n"348 f"Try:\n"349 f"• Rephrasing your question more simply\n"350 f"• Being more specific about what you want to know\n"351 f"• Uploading relevant documents for better context")352 else:353 return (f"I received your question: '{query}'\n\n"354 f"Unfortunately, I cannot provide a detailed answer because:\n"355 f"• The AI model failed to load (likely network connectivity issue)\n"356 f"• This appears to be a connection problem with huggingface.co\n\n"357 f"To resolve this:\n"358 f"• Check your internet connection\n"359 f"• Try again in a few minutes\n"360 f"• Consider using a VPN if there are regional restrictions\n\n"361 f"The app can still search your documents - try uploading PDFs and asking questions about them!")362 363 def _clean_response(self, response: str) -> str:364 """Clean up the generated response"""365 # Remove common artifacts366 response = response.strip()367 368 # Stop at certain tokens that indicate end of response369 stop_tokens = ["\nUser:", "\nBot:", "Question:", "Context:", "Answer:", "<|endoftext|>"]370 for token in stop_tokens:371 if token in response:372 response = response.split(token)[0]373 374 # Remove repetitive patterns (but be more lenient)375 lines = response.split('\n')376 if len(lines) > 1:377 unique_lines = []378 for line in lines:379 line = line.strip()380 if line and line not in unique_lines:381 unique_lines.append(line)382 response = ' '.join(unique_lines)383 384 # Only remove if response is very short (reduced threshold)385 if len(response.strip()) < 3:386 return ""387 388 return response.strip()389 390 def is_available(self) -> bool:391 """Check if the model is available for use"""392 try:393 if not self.is_loaded:394 success = self._load_model()395 return success396 return self.is_loaded397 except Exception as e:398 print(f"Error checking model availability: {str(e)}")399 return False400 401 def get_model_info(self) -> Dict:402 """Get information about the loaded model"""403 return {404 "model_name": self.model_name,405 "device": self.device,406 "is_loaded": self.is_loaded,407 "cache_dir": self.cache_dir,408 "size_mb": self._get_model_size_mb() if self.is_loaded else 0409 }410 411 412class HuggingFaceEmbeddingModel:413 """414 Embedding model using Sentence Transformers with automatic downloading415 """416 417 def __init__(self, model_name: str = None, cache_dir: str = None):418 self.model_name = model_name or config.EMBEDDING_MODEL419 self.cache_dir = cache_dir or config.MODEL_CACHE_DIR420 self.model = None421 self.device = self._setup_device()422 423 # Create cache directory424 os.makedirs(self.cache_dir, exist_ok=True)425 426 print(f"Embedding model: {self.model_name}")427 428 def _setup_device(self):429 """Setup computation device"""430 if config.USE_CUDA and torch.cuda.is_available():431 return "cuda"432 return "cpu"433 434 def _load_model(self):435 """Load the sentence transformer model"""436 if self.model is not None:437 return438 439 try:440 print(f"Loading embedding model: {self.model_name}")441 from sentence_transformers import SentenceTransformer442 443 # Load with explicit device=None to let the library handle device assignment444 self.model = SentenceTransformer(445 self.model_name,446 cache_folder=self.cache_dir,447 device=None, # Let the library choose the best device448 trust_remote_code=True449 )450 451 print(f"Embedding model loaded successfully!")452 except Exception as e:453 print(f"Error loading embedding model: {str(e)}")454 raise e455 456 def encode(self, texts: List[str]) -> torch.Tensor:457 """Encode texts to embeddings"""458 if self.model is None:459 self._load_model()460 461 try:462 embeddings = self.model.encode(texts, convert_to_tensor=True)463 return embeddings.cpu().numpy()464 except Exception as e:465 print(f"Error encoding texts: {str(e)}")466 # Return dummy embeddings as fallback467 import numpy as np468 return np.random.rand(len(texts), 384).astype('float32')469 470 def get_dimension(self) -> int:471 """Get embedding dimension"""472 if self.model is None:473 self._load_model()474 475 # Test with sample text476 sample_embedding = self.encode(["sample text"])477 return sample_embedding.shape[1]478 479 def is_available(self) -> bool:480 """Check if embedding model is available"""481 try:482 if self.model is None:483 self._load_model()484 return self.model is not None485 except:486 return False