Upload 5 files

src/inference/model.py

@@ -1,171 +1,237 @@
-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
-            
-            # 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",
-                trust_remote_code=True,
-                
-            )
-            
-            # Load the LoRA adapter from Hugging Face
-            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",
-                
-            )
-            
-            # Set to evaluation mode
-            self.model.eval()
-            
-            logger.info("Model loaded successfully from Hugging Face")
-            
-        except Exception as e:
-            logger.error(f"Failed to load model: {e}")
-            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']
+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
+            
+<<<<<<< HEAD
+            # 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",
+                    trust_remote_code=True,
+                    attn_implementation="eager",
+                    use_flash_attention_2=False
+                )
+                
+                # 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",
+                        attn_implementation="eager",
+                        use_flash_attention_2=False
+                    )
+                    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
+=======
+            # 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",
+                trust_remote_code=True,
+                
+            )
+            
+            # Load the LoRA adapter from Hugging Face
+            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",
+                
+            )
+>>>>>>> 3b1d9d4700da14631c2d7f96e38c9e460a1a4dd0
+            
+            # Set to evaluation mode
+            self.model.eval()
+            
+<<<<<<< HEAD
+            # 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()}")
+=======
+            logger.info("Model loaded successfully from Hugging Face")
+            
+        except Exception as e:
+            logger.error(f"Failed to load model: {e}")
+>>>>>>> 3b1d9d4700da14631c2d7f96e38c9e460a1a4dd0
+            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']
         } 

src/training/finetune.py

@@ -1,3 +1,289 @@
+<<<<<<< HEAD
+import os
+import sys
+import yaml
+import argparse
+import logging
+from typing import Dict, Any
+import torch
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling,
+    EarlyStoppingCallback,
+    BitsAndBytesConfig
+)
+from peft import (
+    LoraConfig,
+    get_peft_model,
+    prepare_model_for_kbit_training,
+    TaskType
+)
+from datasets import Dataset
+from tqdm import tqdm
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class AgriQAFineTuner:
+    
+    def __init__(self, config_path: str):
+        self.config = self.load_config(config_path) # load the config file
+        self.setup_environment()
+        
+    def load_config(self, config_path: str) -> Dict[str, Any]:
+        with open(config_path, 'r') as f:
+            config = yaml.safe_load(f)
+        return config
+    
+    def setup_environment(self) -> None:
+        # Set device
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"Using device: {self.device}")
+        
+        # Create output directory
+        os.makedirs(self.config['training']['output_dir'], exist_ok=True)
+        
+    def load_model_and_tokenizer(self):
+        logger.info(f"Loading model: {self.config['model']['base_model']}")
+        
+        # Load tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.config['model']['base_model'],
+            trust_remote_code=self.config['model']['trust_remote_code']
+        )
+        
+        # Add padding token if not present
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        
+        # Load model with quantization if specified
+        if self.config['hardware']['use_4bit']:
+            logger.info("Loading model with 4-bit quantization")
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_quant_type=self.config['hardware']['bnb_4bit_quant_type'],
+                bnb_4bit_use_double_quant=self.config['hardware']['bnb_4bit_use_double_quant'],
+            )
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.config['model']['base_model'],
+                quantization_config=quantization_config,
+                device_map=self.config['hardware']['device_map'],
+                trust_remote_code=self.config['model']['trust_remote_code']
+            )
+        else:
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.config['model']['base_model'],
+                device_map=self.config['hardware']['device_map'],
+                trust_remote_code=self.config['model']['trust_remote_code']
+            )
+        
+        # Prepare model for k-bit training
+        if self.config['hardware']['use_4bit']:
+            self.model = prepare_model_for_kbit_training(self.model)
+        
+        logger.info("Model and tokenizer loaded successfully")
+    
+    def setup_lora(self):
+        # Apply LoRA configuration
+        logger.info("Setting up LoRA configuration")
+        lora_config = LoraConfig(
+            r=self.config['lora']['r'],
+            lora_alpha=self.config['lora']['lora_alpha'],
+            target_modules=self.config['lora']['target_modules'],
+            lora_dropout=self.config['lora']['lora_dropout'],
+            bias=self.config['lora']['bias'],
+            task_type=self.config['lora']['task_type'],
+        )
+        
+        # Enable gradient checkpointing for memory optimization
+        if self.config['training']['gradient_checkpointing']:
+            self.model.gradient_checkpointing_enable()
+            logger.info("Gradient checkpointing enabled for memory optimization")
+        
+        # Apply LoRA
+        self.model = get_peft_model(self.model, lora_config)
+        self.model.print_trainable_parameters()
+        
+        logger.info("LoRA configuration applied successfully")
+    
+    def load_dataset(self):
+        """Load the tokenized datasets."""
+        logger.info("Loading dataset")
+        
+        # Load pre-tokenized datasets
+        logger.info("Loading pre-tokenized datasets...")
+        train_dataset = Dataset.load_from_disk(os.path.join(self.config['data']['tokenized_dir'], "train"))
+        val_dataset = Dataset.load_from_disk(os.path.join(self.config['data']['tokenized_dir'], "validation"))
+        
+        # Limit samples if specified
+        max_samples = self.config['data'].get('max_samples', None)
+        if max_samples:
+            logger.info(f"Limiting training samples to {max_samples}")
+            train_dataset = train_dataset.select(range(min(max_samples, len(train_dataset))))
+            val_dataset = val_dataset.select(range(min(max_samples // 10, len(val_dataset))))  # 10% for validation
+        
+        logger.info(f"Loaded tokenized training samples: {len(train_dataset)}")
+        logger.info(f"Loaded tokenized validation samples: {len(val_dataset)}")
+        
+        return train_dataset, val_dataset
+    
+    def setup_training(self, train_dataset, val_dataset):
+        logger.info("Setting up training configuration")
+        
+        # Convert numeric values from config
+        def convert_numeric(value):
+            if isinstance(value, str):
+                try:
+                    return float(value)
+                except ValueError:
+                    return value
+            return value
+        
+        # Training arguments with memory optimizations
+        training_args = TrainingArguments(
+            output_dir=self.config['training']['output_dir'],
+            num_train_epochs=convert_numeric(self.config['training']['num_train_epochs']),
+            per_device_train_batch_size=convert_numeric(self.config['training']['per_device_train_batch_size']),
+            per_device_eval_batch_size=convert_numeric(self.config['training']['per_device_eval_batch_size']),
+            gradient_accumulation_steps=convert_numeric(self.config['training']['gradient_accumulation_steps']),
+            learning_rate=convert_numeric(self.config['training']['learning_rate']),
+            weight_decay=convert_numeric(self.config['training']['weight_decay']),
+            warmup_steps=convert_numeric(self.config['training']['warmup_steps']),
+            logging_steps=convert_numeric(self.config['training']['logging_steps']),
+            save_steps=convert_numeric(self.config['training']['save_steps']),
+            eval_steps=convert_numeric(self.config['training']['eval_steps']),
+            evaluation_strategy=self.config['training']['evaluation_strategy'],
+            save_strategy=self.config['training']['save_strategy'],
+            save_total_limit=convert_numeric(self.config['training']['save_total_limit']),
+            load_best_model_at_end=self.config['training']['load_best_model_at_end'],
+            metric_for_best_model=self.config['training']['metric_for_best_model'],
+            greater_is_better=self.config['training']['greater_is_better'],
+            fp16=self.config['training']['fp16'],
+            dataloader_num_workers=convert_numeric(self.config['training']['dataloader_num_workers']),
+            gradient_checkpointing=self.config['training']['gradient_checkpointing'],
+            max_grad_norm=convert_numeric(self.config['training']['max_grad_norm']),
+            report_to=self.config['logging']['report_to'],
+            run_name=self.config['logging']['run_name'],
+            log_level=self.config['logging']['log_level'],
+            # Memory optimization settings
+            dataloader_drop_last=True,
+            group_by_length=True,
+            length_column_name="length",
+            # Disable features that use more memory
+            ddp_find_unused_parameters=False,
+            dataloader_pin_memory=False,
+            # Additional memory optimizations
+            optim="adamw_torch_fused",  # Use fused optimizer for speed
+            torch_compile=False,  # Disable torch.compile for memory
+            use_cpu=False,  # Keep on GPU but optimize memory
+            # Reduce memory fragmentation
+            dataloader_persistent_workers=False,
+        )
+        
+        # Data collator for pre-tokenized data
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=self.tokenizer,
+            mlm=False,
+        )
+        
+        # Trainer
+        self.trainer = Trainer(
+            model=self.model,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=val_dataset,
+            data_collator=data_collator,
+            callbacks=[EarlyStoppingCallback(early_stopping_patience=3)]
+        )
+        
+        logger.info("Training setup completed")
+    
+    def train(self):
+        logger.info("Starting training...")
+        
+        try:
+            # Train the model
+            train_result = self.trainer.train()
+            
+            # Save the final model
+            self.trainer.save_model()
+            
+            # Save training metrics
+            metrics = train_result.metrics
+            self.trainer.log_metrics("train", metrics)
+            self.trainer.save_metrics("train", metrics)
+            self.trainer.save_state()
+            
+            logger.info("Training completed successfully!")
+            logger.info(f"Training metrics: {metrics}")
+            
+        except Exception as e:
+            logger.error(f"Training failed: {e}")
+            raise
+    
+    def save_model(self):
+        logger.info("Saving model...")
+        
+        output_dir = self.config['training']['output_dir']
+        
+        # Save tokenizer
+        self.tokenizer.save_pretrained(output_dir)
+        
+        # Save model configuration
+        model_config = {
+            'base_model': self.config['model']['base_model'],
+            'lora_config': self.config['lora'],
+            'generation_config': self.config['generation']
+        }
+        
+        config_path = os.path.join(output_dir, 'model_config.json')
+        import json
+        with open(config_path, 'w') as f:
+            json.dump(model_config, f, indent=2)
+        
+        logger.info(f"Model saved to {output_dir}")
+    
+    def run(self):
+        logger.info("Starting agriQA fine-tuning pipeline...")
+        
+        # Load model and tokenizer
+        self.load_model_and_tokenizer()
+        
+        # Setup LoRA
+        self.setup_lora()
+        
+        # Load and prepare datasets
+        train_dataset, val_dataset = self.load_dataset()
+        
+        # Setup training
+        self.setup_training(train_dataset, val_dataset)
+        
+        # Train the model
+        self.train()
+        
+        # Save the model
+        self.save_model()
+        
+        logger.info("Fine-tuning pipeline completed successfully!")
+
+def main():
+    parser = argparse.ArgumentParser(description="Fine-tune Qwen model on agriQA dataset")
+    parser.add_argument("--config", type=str, default="configs/training_config.yaml",
+                       help="Path to training configuration file")
+    
+    args = parser.parse_args()
+    
+    # Initialize and run fine-tuning
+    fine_tuner = AgriQAFineTuner(args.config)
+    fine_tuner.run()
+
+if __name__ == "__main__":
+=======
 import os
 import sys
 import yaml
@@ -282,4 +568,5 @@ def main():
     fine_tuner.run()
 
 if __name__ == "__main__":
+>>>>>>> 3b1d9d4700da14631c2d7f96e38c9e460a1a4dd0
     main()