Create training_manager.py

training_manager.py

@@ -0,0 +1,174 @@
+"""
+PyPilot Training Manager - Advanced distributed training with monitoring
+"""
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, Dataset
+from transformers import TrainingArguments, Trainer, EarlyStoppingCallback
+import wandb
+import numpy as np
+import time
+from datetime import datetime
+import os
+
+class CodeDataset(Dataset):
+    def __init__(self, tokenized_data):
+        self.data = tokenized_data
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        return self.data[idx]
+
+class PyPilotTrainingManager:
+    def __init__(self, model, model_name="PyPilot"):
+        self.model = model
+        self.model_name = model_name
+        self.training_history = []
+        self.best_loss = float('inf')
+        
+    def setup_distributed_training(self, use_fp16=True, use_gradient_checkpointing=True):
+        """Configure distributed training options"""
+        training_args = TrainingArguments(
+            output_dir=f"./pypilot-checkpoints",
+            overwrite_output_dir=True,
+            num_train_epochs=10,
+            per_device_train_batch_size=4,
+            per_device_eval_batch_size=4,
+            gradient_accumulation_steps=8,
+            learning_rate=5e-5,
+            weight_decay=0.01,
+            warmup_steps=1000,
+            logging_dir="./logs",
+            logging_steps=500,
+            eval_steps=1000,
+            save_steps=2000,
+            save_total_limit=5,
+            prediction_loss_only=True,
+            remove_unused_columns=False,
+            fp16=use_fp16,
+            dataloader_pin_memory=False,
+            gradient_checkpointing=use_gradient_checkpointing,
+            report_to=["wandb"],
+            run_name=f"pypilot-{datetime.now().strftime('%Y-%m-%d-%H-%M')}",
+        )
+        return training_args
+    
+    def setup_wandb_monitoring(self, project_name="pypilot"):
+        """Initialize Weights & Biases for experiment tracking"""
+        wandb.init(
+            project=project_name,
+            name=f"pypilot-{datetime.now().strftime('%Y-%m-%d-%H-%M')}",
+            config={
+                "architecture": "Transformer",
+                "dataset": "GitHub Code",
+                "epochs": 10,
+                "batch_size": 32,
+            }
+        )
+    
+    def create_advanced_callbacks(self):
+        """Create callbacks for training optimization"""
+        callbacks = [
+            EarlyStoppingCallback(early_stopping_patience=3),
+        ]
+        return callbacks
+    
+    def compute_metrics(self, eval_pred):
+        """Compute advanced metrics for code generation"""
+        predictions, labels = eval_pred
+        predictions = torch.tensor(predictions)
+        labels = torch.tensor(labels)
+        
+        # Calculate perplexity
+        loss_fct = nn.CrossEntropyLoss()
+        loss = loss_fct(predictions.view(-1, predictions.size(-1)), labels.view(-1))
+        perplexity = torch.exp(loss)
+        
+        # Calculate accuracy
+        preds = torch.argmax(predictions, dim=-1)
+        accuracy = (preds == labels).float().mean()
+        
+        return {
+            "perplexity": perplexity.item(),
+            "accuracy": accuracy.item(),
+            "loss": loss.item()
+        }
+    
+    def train_with_advanced_features(self, train_dataset, eval_dataset=None):
+        """Start advanced training with all features"""
+        print("🚀 Starting Advanced PyPilot Training...")
+        
+        # Setup monitoring
+        self.setup_wandb_monitoring()
+        
+        # Configure training
+        training_args = self.setup_distributed_training()
+        callbacks = self.create_advanced_callbacks()
+        
+        # Create trainer
+        trainer = Trainer(
+            model=self.model,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            compute_metrics=self.compute_metrics,
+            callbacks=callbacks,
+        )
+        
+        # Start training
+        print("🎯 Training started with advanced features:")
+        print(f"   - FP16 Precision: Enabled")
+        print(f"   - Gradient Checkpointing: Enabled") 
+        print(f"   - Early Stopping: Enabled")
+        print(f"   - W&B Monitoring: Enabled")
+        
+        trainer.train()
+        
+        # Save final model
+        trainer.save_model("./pypilot-final-model")
+        print("✅ Training completed and model saved!")
+        
+        return trainer
+    
+    def hyperparameter_search(self, train_dataset, param_combinations):
+        """Perform hyperparameter search"""
+        best_params = None
+        
+        for i, params in enumerate(param_combinations):
+            print(f"🔍 Testing hyperparameter combination {i+1}/{len(param_combinations)}")
+            
+            # Update model with new params
+            self.update_model_hyperparams(params)
+            
+            # Quick training run to evaluate
+            quick_trainer = Trainer(
+                model=self.model,
+                args=TrainingArguments(
+                    output_dir=f"./hparam-search-{i}",
+                    num_train_epochs=1,
+                    per_device_train_batch_size=params['batch_size'],
+                    learning_rate=params['learning_rate'],
+                ),
+                train_dataset=train_dataset,
+            )
+            
+            results = quick_trainer.train()
+            
+            if results.training_loss < self.best_loss:
+                self.best_loss = results.training_loss
+                best_params = params
+        
+        print(f"🎯 Best hyperparameters: {best_params}")
+        return best_params
+
+if __name__ == "__main__":
+    # Example usage
+    from modeling_pypilot import PyPilotModel, PyPilotConfig
+    
+    config = PyPilotConfig()
+    model = PyPilotModel(config)
+    
+    manager = PyPilotTrainingManager(model)
+    print("✅ Training Manager ready!")