Upload .\src\training\trainer.py with huggingface_hub

.//src//training//trainer.py

@@ -0,0 +1,230 @@
+"""Training pipeline for BwengeAi."""
+
+import logging
+import os
+from pathlib import Path
+from typing import Any
+
+import torch
+from datasets import load_dataset
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    DataCollatorForLanguageModeling,
+    TrainingArguments,
+)
+from trl import SFTTrainer
+
+logger = logging.getLogger(__name__)
+
+
+class BwengeTrainer:
+    """Training manager for BwengeAi."""
+
+    def __init__(self, config: dict[str, Any]):
+        self.config = config
+        self.training_config = config.get("training", {})
+        self.output_dir = Path(self.training_config.get("output_dir", "outputs"))
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+    def prepare_dataset(
+        self,
+        data_path: str,
+        tokenizer: AutoTokenizer,
+        max_length: int = 2048,
+    ) -> Any:
+        """Prepare dataset for training."""
+        logger.info(f"Loading dataset from {data_path}")
+        
+        dataset = load_dataset("json", data_files=data_path, split="train")
+        
+        def tokenize_function(examples):
+            return tokenizer(
+                examples["text"],
+                truncation=True,
+                max_length=max_length,
+                padding="max_length",
+            )
+        
+        tokenized_dataset = dataset.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=dataset.column_names,
+        )
+        
+        logger.info(f"Dataset prepared: {len(tokenized_dataset)} samples")
+        return tokenized_dataset
+
+    def prepare_instruction_dataset(
+        self,
+        data_path: str,
+        tokenizer: AutoTokenizer,
+        max_length: int = 2048,
+    ) -> Any:
+        """Prepare instruction-following dataset."""
+        logger.info(f"Loading instruction dataset from {data_path}")
+        
+        dataset = load_dataset("json", data_files=data_path, split="train")
+        
+        def format_instruction(examples):
+            texts = []
+            for i in range(len(examples["instruction"])):
+                instruction = examples["instruction"][i]
+                input_text = examples.get("input", [""] * len(examples["instruction"]))[i]
+                output = examples["output"][i]
+                
+                if input_text:
+                    text = f"### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:\n{output}"
+                else:
+                    text = f"### Instruction:\n{instruction}\n\n### Response:\n{output}"
+                
+                texts.append(text)
+            
+            return {"text": texts}
+        
+        formatted_dataset = dataset.map(
+            format_instruction,
+            batched=True,
+            remove_columns=dataset.column_names,
+        )
+        
+        def tokenize_function(examples):
+            return tokenizer(
+                examples["text"],
+                truncation=True,
+                max_length=max_length,
+                padding="max_length",
+            )
+        
+        tokenized_dataset = formatted_dataset.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=formatted_dataset.column_names,
+        )
+        
+        logger.info(f"Instruction dataset prepared: {len(tokenized_dataset)} samples")
+        return tokenized_dataset
+
+    def train(
+        self,
+        model: AutoModelForCausalLM,
+        tokenizer: AutoTokenizer,
+        dataset: Any,
+        lora: bool = False,
+    ) -> None:
+        """Train the model."""
+        training_args = TrainingArguments(
+            output_dir=str(self.output_dir),
+            num_train_epochs=self.training_config.get("num_epochs", 3),
+            per_device_train_batch_size=self.training_config.get("batch_size", 8),
+            gradient_accumulation_steps=self.training_config.get("gradient_accumulation_steps", 4),
+            learning_rate=self.training_config.get("learning_rate", 2e-5),
+            weight_decay=self.training_config.get("weight_decay", 0.01),
+            warmup_steps=self.training_config.get("warmup_steps", 500),
+            max_grad_norm=self.training_config.get("max_grad_norm", 1.0),
+            fp16=self.training_config.get("fp16", False) and torch.cuda.is_available(),
+            logging_steps=self.training_config.get("logging_steps", 10),
+            save_steps=self.training_config.get("save_steps", 500),
+            save_total_limit=self.training_config.get("save_total_limit", 3),
+            report_to="none",
+            seed=self.training_config.get("seed", 42),
+            dataloader_num_workers=0,
+            remove_unused_columns=False,
+        )
+        
+        if lora:
+            trainer = SFTTrainer(
+                model=model,
+                train_dataset=dataset,
+                args=training_args,
+                tokenizer=tokenizer,
+                max_seq_length=self.config.get("model", {}).get("max_length", 2048),
+            )
+        else:
+            data_collator = DataCollatorForLanguageModeling(
+                tokenizer=tokenizer,
+                mlm=False,
+            )
+            
+            trainer = SFTTrainer(
+                model=model,
+                train_dataset=dataset,
+                args=training_args,
+                tokenizer=tokenizer,
+                data_collator=data_collator,
+                max_seq_length=self.config.get("model", {}).get("max_length", 2048),
+            )
+        
+        logger.info("Starting training...")
+        trainer.train()
+        
+        trainer.save_model(str(self.output_dir / "final"))
+        tokenizer.save_pretrained(str(self.output_dir / "final"))
+        
+        logger.info(f"Training complete. Model saved to {self.output_dir / 'final'}")
+
+    def train_from_config(
+        self,
+        model: AutoModelForCausalLM,
+        tokenizer: AutoTokenizer,
+        data_path: str,
+        lora: bool = False,
+    ) -> None:
+        """Train using configuration."""
+        logger.info(f"Loading dataset from {data_path}")
+        dataset = load_dataset("json", data_files=data_path, split="train")
+        
+        has_text = "text" in dataset.column_names
+        has_instruction = "instruction" in dataset.column_names
+        
+        if has_instruction and not has_text:
+            def format_instruction(examples):
+                texts = []
+                for i in range(len(examples["instruction"])):
+                    instruction = examples["instruction"][i]
+                    input_text = examples.get("input", [""] * len(examples["instruction"]))[i]
+                    output = examples["output"][i]
+                    if input_text:
+                        text = f"### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:\n{output}"
+                    else:
+                        text = f"### Instruction:\n{instruction}\n\n### Response:\n{output}"
+                    texts.append(text)
+                return {"text": texts}
+            
+            dataset = dataset.map(format_instruction, batched=True, remove_columns=dataset.column_names)
+        
+        max_seq_length = self.config.get("model", {}).get("max_length", 2048)
+        
+        training_args = TrainingArguments(
+            output_dir=str(self.output_dir),
+            num_train_epochs=self.training_config.get("num_epochs", 3),
+            per_device_train_batch_size=self.training_config.get("batch_size", 8),
+            gradient_accumulation_steps=self.training_config.get("gradient_accumulation_steps", 4),
+            learning_rate=self.training_config.get("learning_rate", 2e-5),
+            weight_decay=self.training_config.get("weight_decay", 0.01),
+            warmup_steps=self.training_config.get("warmup_steps", 500),
+            max_grad_norm=self.training_config.get("max_grad_norm", 1.0),
+            fp16=self.training_config.get("fp16", False) and torch.cuda.is_available(),
+            logging_steps=self.training_config.get("logging_steps", 10),
+            save_steps=self.training_config.get("save_steps", 500),
+            save_total_limit=self.training_config.get("save_total_limit", 3),
+            report_to="none",
+            seed=self.training_config.get("seed", 42),
+            dataloader_num_workers=0,
+            remove_unused_columns=False,
+        )
+        
+        trainer = SFTTrainer(
+            model=model,
+            train_dataset=dataset,
+            args=training_args,
+            processing_class=tokenizer,
+        )
+        
+        logger.info("Starting training...")
+        trainer.train()
+        
+        trainer.save_model(str(self.output_dir / "final"))
+        tokenizer.save_pretrained(str(self.output_dir / "final"))
+        
+        logger.info(f"Training complete. Model saved to {self.output_dir / 'final'}")