n

src/inference/model.py

@@ -1,23 +1,15 @@
 import os
 import json
-import logging
-
-# Set PyTorch environment variables before import to prevent logging issues
-os.environ.setdefault('TORCH_LOGS', 'torch')
-os.environ.setdefault('TORCH_SHOW_CPP_STACKTRACES', '0')
-os.environ.setdefault('TORCH_USE_CUDA_DSA', '0')
-
 import torch
+import logging
 from typing import Dict, Any, Optional
-
-# Set transformers environment variables
-os.environ.setdefault('TRANSFORMERS_VERBOSITY', 'error')
-os.environ.setdefault('TOKENIZERS_PARALLELISM', 'false')
-
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
 import time
 
+# Fix PyTorch logging issue
+os.environ['TORCH_LOGS'] = 'torch'
+
 logger = logging.getLogger(__name__)
 
 class AgriQAAssistant:
@@ -32,99 +24,85 @@ class AgriQAAssistant:
         self.load_model()
     
     def load_model(self):
-    
         logger.info(f"Loading model from Hugging Face: {self.model_path}")
         
-        # Set additional environment variables for model loading
-        os.environ.setdefault('HF_HUB_OFFLINE', 'false')
-        os.environ.setdefault('HF_HUB_DISABLE_TELEMETRY', '1')
+        # Configuration for the uploaded model
+        self.config = {
+            'base_model': 'Qwen/Qwen1.5-1.8B-Chat',
+            'generation_config': {
+                'max_new_tokens': 512,
+                'do_sample': True,
+                'temperature': 0.3,
+                'top_p': 0.85,
+                'top_k': 40,
+                'repetition_penalty': 1.2,
+                'length_penalty': 1.1,
+                'no_repeat_ngram_size': 3
+            }
+        }
+        
+        # Load tokenizer from base model
+        logger.info("Loading tokenizer from base model...")
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.config['base_model'],
+            trust_remote_code=True
+        )
+        
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
         
+        # Try to load the model directly from Hugging Face first
         try:
-            # Configuration for the uploaded model
-            self.config = {
-                'base_model': 'Qwen/Qwen1.5-1.8B-Chat',
-                'generation_config': {
-                    'max_new_tokens': 512,  # Increased for complete responses
-                    'do_sample': True,
-                    'temperature': 0.3,     # Lower temperature for more consistent, structured responses
-                    'top_p': 0.85,         # Slightly lower for more focused sampling
-                    'top_k': 40,           # Lower for more focused responses
-                    'repetition_penalty': 1.2,  # Higher penalty to avoid repetition
-                    'length_penalty': 1.1,      # Encourage slightly longer, detailed responses
-                    'no_repeat_ngram_size': 3   # Avoid repeating 3-grams
-                }
-            }
+            logger.info("Attempting to load model directly from Hugging Face...")
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_path,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                trust_remote_code=True,
+                attn_implementation="eager",
+                use_flash_attention_2=False
+            )
+            logger.info("Model loaded directly from Hugging Face successfully")
+        except Exception as direct_load_error:
+            logger.info(f"Direct loading failed: {direct_load_error}")
+            logger.info("Falling back to base model + LoRA adapter approach...")
             
-            # Load tokenizer from base model
-            logger.info("Loading tokenizer from base model...")
-            self.tokenizer = AutoTokenizer.from_pretrained(
+            # Load base model first
+            logger.info("Loading base model...")
+            base_model = AutoModelForCausalLM.from_pretrained(
                 self.config['base_model'],
-                trust_remote_code=True
+                torch_dtype=torch.float16,
+                device_map="auto"
             )
             
-            if self.tokenizer.pad_token is None:
-                self.tokenizer.pad_token = self.tokenizer.eos_token
-            
-            # Try to load the model directly from Hugging Face first
+            # Try to load the LoRA adapter
             try:
-                logger.info("Attempting to load model directly from Hugging Face...")
-                self.model = AutoModelForCausalLM.from_pretrained(
+                logger.info("Loading LoRA adapter from Hugging Face...")
+                self.model = PeftModel.from_pretrained(
+                    base_model,
                     self.model_path,
                     torch_dtype=torch.float16,
-                    device_map="auto",
-                    trust_remote_code=True,
-                    attn_implementation="eager",
-                    use_flash_attention_2=False
-                )
-                logger.info("Model loaded directly from Hugging Face successfully")
-            except Exception as direct_load_error:
-                logger.info(f"Direct loading failed: {direct_load_error}")
-                logger.info("Falling back to base model + LoRA adapter approach...")
-                
-                # Load base model first
-                logger.info("Loading base model...")
-                base_model = AutoModelForCausalLM.from_pretrained(
-                    self.config['base_model'],
-                    torch_dtype=torch.float16,
                     device_map="auto"
                 )
-                
-                # Try to load the LoRA adapter
-                try:
-                    logger.info("Loading LoRA adapter from Hugging Face...")
-                    self.model = PeftModel.from_pretrained(
-                        base_model,
-                        self.model_path,
-                        torch_dtype=torch.float16,
-                        device_map="auto"
-                    )
-                    logger.info("LoRA adapter loaded successfully")
-                except Exception as lora_error:
-                    logger.warning(f"LoRA adapter loading failed: {lora_error}")
-                    logger.info("Using base model without LoRA adapter...")
-                    self.model = base_model
-            
-            # Set to evaluation mode
-            self.model.eval()
-            
-            # Log model information
-            logger.info(f"Model loaded successfully from Hugging Face")
-            logger.info(f"Model type: {type(self.model).__name__}")
-            logger.info(f"Device: {self.device}")
-            
-            # Check if it's a PeftModel
-            if hasattr(self.model, 'peft_config'):
-                logger.info("LoRA adapter configuration:")
-                for adapter_name, config in self.model.peft_config.items():
-                    logger.info(f"  - {adapter_name}: {config.target_modules}")
-            
-        except Exception as e:
-            logger.error(f"Failed to load model: {e}")
-            logger.error(f"Model path: {self.model_path}")
-            logger.error(f"Base model: {self.config['base_model']}")
-            import traceback
-            logger.error(f"Traceback: {traceback.format_exc()}")
-            raise
+                logger.info("LoRA adapter loaded successfully")
+            except Exception as lora_error:
+                logger.warning(f"LoRA adapter loading failed: {lora_error}")
+                logger.info("Using base model without LoRA adapter...")
+                self.model = base_model
+        
+        # Set to evaluation mode
+        self.model.eval()
+        
+        # Log model information
+        logger.info(f"Model loaded successfully from Hugging Face")
+        logger.info(f"Model type: {type(self.model).__name__}")
+        logger.info(f"Device: {self.device}")
+        
+        # Check if it's a PeftModel
+        if hasattr(self.model, 'peft_config'):
+            logger.info("LoRA adapter configuration:")
+            for adapter_name, config in self.model.peft_config.items():
+                logger.info(f"  - {adapter_name}: {config.target_modules}")
     
     def format_prompt(self, question: str) -> str:
         """Format the question for the model using proper format."""
@@ -152,72 +130,49 @@ class AgriQAAssistant:
     def generate_response(self, question: str, max_length: Optional[int] = None) -> Dict[str, Any]:
         start_time = time.time()
         
-        try:
-            # Format the prompt
-            prompt = self.format_prompt(question)
-            
-            # Set device for inputs
-            device = self.device if hasattr(self, 'device') else 'cpu'
-            
-            # Tokenize input
-            inputs = self.tokenizer(
-                prompt,
-                return_tensors="pt",
-                truncation=True,
-                max_length=2048
-            ).to(device)
-            
-            # Generation parameters
-            gen_config = self.config['generation_config'].copy()
-            if max_length:
-                gen_config['max_new_tokens'] = max_length
-            
-            # Generate response
-            with torch.no_grad():
-                try:
-                    outputs = self.model.generate(
-                        **inputs,
-                        **gen_config,
-                        pad_token_id=self.tokenizer.eos_token_id
-                    )
-                except Exception as gen_error:
-                    logger.error(f"Generation error: {gen_error}")
-                    # Fallback to simpler generation
-                    outputs = self.model.generate(
-                        **inputs,
-                        max_new_tokens=gen_config.get('max_new_tokens', 512),
-                        do_sample=False,
-                        pad_token_id=self.tokenizer.eos_token_id
-                    )
-            
-            # Decode response
-            response = self.tokenizer.decode(
-                outputs[0][inputs['input_ids'].shape[1]:],
-                skip_special_tokens=True
-            ).strip()
-            
-            # Calculate response time
-            response_time = time.time() - start_time
-            
-            return {
-                'answer': response,
-                'response_time': response_time,
-                'model_info': {
-                    'model_name': 'agriqa-assistant',
-                    'model_source': 'Hugging Face',
-                    'model_path': self.model_path,
-                    'base_model': self.config['base_model']
-                }
-            }
-            
-        except Exception as e:
-            logger.error(f"Error generating response: {e}")
-            return {
-                'answer': "I apologize, but I encountered an error while processing your question. Please try again.",
-                'confidence': 0.0,
-                'response_time': time.time() - start_time,
-                'error': str(e)
+        # Format the prompt
+        prompt = self.format_prompt(question)
+        
+        # Tokenize input
+        inputs = self.tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=2048
+        ).to(self.device)
+        
+        # Generation parameters
+        gen_config = self.config['generation_config'].copy()
+        if max_length:
+            gen_config['max_new_tokens'] = max_length
+        
+        # Generate response
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                **gen_config,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
+        
+        # Decode response
+        response = self.tokenizer.decode(
+            outputs[0][inputs['input_ids'].shape[1]:],
+            skip_special_tokens=True
+        ).strip()
+        
+        # Calculate response time
+        response_time = time.time() - start_time
+        
+        return {
+            'answer': response,
+            'response_time': response_time,
+            'model_info': {
+                'model_name': 'agriqa-assistant',
+                'model_source': 'Hugging Face',
+                'model_path': self.model_path,
+                'base_model': self.config['base_model']
             }
+        }
     
     def get_model_info(self) -> Dict[str, Any]:
         """Get information about the loaded model."""