Upload model.py

model.py

@@ -0,0 +1,203 @@
+import os
+import json
+import torch
+import logging
+from typing import Dict, Any, Optional
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+import time
+
+logger = logging.getLogger(__name__)
+
+class AgriQAAssistant:
+    
+    def __init__(self, model_path: str = "nada013/agriqa-assistant"):
+        self.model_path = model_path
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = None
+        self.tokenizer = None
+        self.config = None
+        
+        self.load_model()
+    
+    def load_model(self):
+    
+        logger.info(f"Loading model from Hugging Face: {self.model_path}")
+        
+        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
+                }
+            }
+            
+            # 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:
+                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 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
+    
+    def format_prompt(self, question: str) -> str:
+        """Format the question for the model using proper format."""
+        # Use the tokenizer's chat template if available
+        if hasattr(self.tokenizer, 'apply_chat_template'):
+            try:
+                messages = [
+                    {"role": "system", "content": "You are AgriQA, an agricultural expert assistant. Your job is to answer farmers' questions with clear, practical, and accurate steps they can directly apply in the field.\n\nWhen answering:\n1. Start with a short, direct answer to the question.\n2. Provide a numbered step-by-step solution.\n3. Include specific details like measurements, quantities, time intervals, and names of products or tools.\n4. Mention any safety precautions if needed.\n5. End with an extra tip or follow-up advice.\n\nFormat Example:\nQuestion: How to control aphid infestation in mustard crops?\nAnswer:\n1. Inspect the crop daily to detect early signs of infestation.\n2. Spray Imidacloprid 17.8% SL at a rate of 0.3 ml per liter of water.\n3. Ensure thorough coverage, especially under the leaves.\n4. Remove surrounding weeds that may host aphids.\n5. Repeat spraying after 7 days if infestation continues.\nNote: Wear gloves and a mask during spraying.\n\nAlways keep your language clear, concise, and easy to understand."},
+                    {"role": "user", "content": question}
+                ]
+                formatted_prompt = self.tokenizer.apply_chat_template(
+                    messages,
+                    tokenize=False,
+                    add_generation_prompt=True
+                )
+                return formatted_prompt
+            except Exception as e:
+                logger.warning(f"Failed to use chat template: {e}. Using fallback format.")
+        
+        # Fallback format for Qwen1.5-Chat
+        system_prompt = "You are AgriQA, an agricultural expert assistant. Your job is to answer farmers' questions with clear, practical, and accurate steps they can directly apply in the field.\n\nWhen answering:\n1. Start with a short, direct answer to the question.\n2. Provide a numbered step-by-step solution.\n3. Include specific details like measurements, quantities, time intervals, and names of products or tools.\n4. Mention any safety precautions if needed.\n5. End with an extra tip or follow-up advice.\n\nFormat Example:\nQuestion: How to control aphid infestation in mustard crops?\nAnswer:\n1. Inspect the crop daily to detect early signs of infestation.\n2. Spray Imidacloprid 17.8% SL at a rate of 0.3 ml per liter of water.\n3. Ensure thorough coverage, especially under the leaves.\n4. Remove surrounding weeds that may host aphids.\n5. Repeat spraying after 7 days if infestation continues.\nNote: Wear gloves and a mask during spraying.\n\nAlways keep your language clear, concise, and easy to understand."
+        formatted_prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant\n"
+        return formatted_prompt
+    
+    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)
+            
+            # 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']
+                }
+            }
+            
+        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)
+            }
+    
+    def get_model_info(self) -> Dict[str, Any]:
+        """Get information about the loaded model."""
+        return {
+            'model_name': 'agriqa-assistant',
+            'model_source': 'Hugging Face',
+            'model_path': self.model_path,
+            'base_model': self.config['base_model'],
+            'device': self.device,
+            'generation_config': self.config['generation_config']
+        }