Create autotrain_v15.py

autotrain_v15.py

@@ -0,0 +1,379 @@
+"""
+Helion-V1.5 AutoTrain Script
+Enhanced training with better error handling and AutoTrain compatibility
+"""
+
+import os
+import sys
+import json
+import logging
+import traceback
+from pathlib import Path
+from dataclasses import dataclass
+from typing import Optional, Dict
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler('helion_v15_training.log'),
+        logging.StreamHandler(sys.stdout)
+    ]
+)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class HelionV15Config:
+    """Configuration for Helion-V1.5 training."""
+    model_name: str = "DeepXR/Helion-V1.5"
+    base_model: str = "meta-llama/Llama-2-7b-hf"
+    dataset_name: str = None
+    output_dir: str = "./helion-v1.5-output"
+    hub_model_id: str = "DeepXR/Helion-V1.5"
+    
+    # Training params
+    num_epochs: int = 3
+    batch_size: int = 4
+    gradient_accumulation: int = 8
+    learning_rate: float = 2e-5
+    warmup_steps: int = 100
+    max_seq_length: int = 4096
+    
+    # LoRA config
+    lora_r: int = 64
+    lora_alpha: int = 128
+    lora_dropout: float = 0.05
+    
+    # AutoTrain specific
+    use_autotrain: bool = True
+    autotrain_backend: str = "local"  # or "spaces"
+    
+    # HuggingFace token
+    hf_token: Optional[str] = None
+
+
+class HelionV15Trainer:
+    """Enhanced trainer for Helion-V1.5 with AutoTrain support."""
+    
+    def __init__(self, config: HelionV15Config):
+        self.config = config
+        self.hf_token = config.hf_token or os.getenv("HF_TOKEN")
+        
+        if not self.hf_token:
+            raise ValueError("HuggingFace token required. Set HF_TOKEN environment variable.")
+    
+    def verify_setup(self) -> bool:
+        """Verify all prerequisites."""
+        logger.info("Verifying setup for Helion-V1.5...")
+        
+        checks = {
+            "CUDA Available": self._check_cuda(),
+            "HuggingFace Token": self._check_token(),
+            "Base Model Access": self._check_base_model(),
+            "Disk Space": self._check_disk_space()
+        }
+        
+        for check, result in checks.items():
+            status = "✅" if result else "❌"
+            logger.info(f"{status} {check}")
+        
+        return all(checks.values())
+    
+    def _check_cuda(self) -> bool:
+        """Check CUDA availability."""
+        try:
+            import torch
+            if torch.cuda.is_available():
+                logger.info(f"Found {torch.cuda.device_count()} GPU(s)")
+                for i in range(torch.cuda.device_count()):
+                    logger.info(f"  GPU {i}: {torch.cuda.get_device_name(i)}")
+                return True
+            return False
+        except:
+            return False
+    
+    def _check_token(self) -> bool:
+        """Verify HuggingFace token."""
+        try:
+            from huggingface_hub import HfApi
+            api = HfApi(token=self.hf_token)
+            user_info = api.whoami()
+            logger.info(f"Logged in as: {user_info['name']}")
+            return True
+        except Exception as e:
+            logger.error(f"Token validation failed: {e}")
+            return False
+    
+    def _check_base_model(self) -> bool:
+        """Check base model access."""
+        try:
+            from huggingface_hub import HfApi
+            api = HfApi(token=self.hf_token)
+            api.model_info(self.config.base_model)
+            return True
+        except Exception as e:
+            logger.error(f"Cannot access base model: {e}")
+            return False
+    
+    def _check_disk_space(self, required_gb: int = 50) -> bool:
+        """Check available disk space."""
+        try:
+            import shutil
+            stat = shutil.disk_usage(self.config.output_dir)
+            available_gb = stat.free / (1024 ** 3)
+            logger.info(f"Available disk space: {available_gb:.2f} GB")
+            return available_gb >= required_gb
+        except:
+            return False
+    
+    def prepare_model(self):
+        """Load and prepare model for training."""
+        import torch
+        from transformers import (
+            AutoTokenizer,
+            AutoModelForCausalLM,
+            BitsAndBytesConfig
+        )
+        from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+        
+        logger.info("Loading tokenizer...")
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.config.base_model,
+            token=self.hf_token,
+            trust_remote_code=True
+        )
+        
+        # Add special tokens
+        special_tokens = {
+            "additional_special_tokens": ["<|system|>", "<|user|>", "<|assistant|>"]
+        }
+        self.tokenizer.add_special_tokens(special_tokens)
+        
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.unk_token
+        
+        logger.info("Loading base model with quantization...")
+        
+        # QLoRA quantization config
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16
+        )
+        
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.config.base_model,
+            quantization_config=bnb_config,
+            device_map="auto",
+            token=self.hf_token,
+            trust_remote_code=True
+        )
+        
+        # Resize embeddings for new tokens
+        self.model.resize_token_embeddings(len(self.tokenizer))
+        
+        # Prepare for k-bit training
+        self.model = prepare_model_for_kbit_training(self.model)
+        
+        # LoRA configuration
+        peft_config = LoraConfig(
+            r=self.config.lora_r,
+            lora_alpha=self.config.lora_alpha,
+            lora_dropout=self.config.lora_dropout,
+            bias="none",
+            task_type="CAUSAL_LM",
+            target_modules=[
+                "q_proj", "k_proj", "v_proj", "o_proj",
+                "gate_proj", "up_proj", "down_proj"
+            ]
+        )
+        
+        self.model = get_peft_model(self.model, peft_config)
+        self.model.print_trainable_parameters()
+        
+        logger.info("✅ Model prepared successfully")
+    
+    def load_dataset(self):
+        """Load training dataset."""
+        from datasets import load_dataset
+        
+        logger.info(f"Loading dataset: {self.config.dataset_name}")
+        
+        self.dataset = load_dataset(
+            self.config.dataset_name,
+            token=self.hf_token
+        )
+        
+        logger.info(f"Dataset loaded: {self.dataset}")
+        
+        # Preprocessing function
+        def preprocess(examples):
+            texts = examples.get("text", [])
+            model_inputs = self.tokenizer(
+                texts,
+                max_length=self.config.max_seq_length,
+                truncation=True,
+                padding="max_length"
+            )
+            model_inputs["labels"] = model_inputs["input_ids"].copy()
+            return model_inputs
+        
+        logger.info("Preprocessing dataset...")
+        self.tokenized_dataset = self.dataset.map(
+            preprocess,
+            batched=True,
+            remove_columns=self.dataset["train"].column_names
+        )
+        
+        logger.info("✅ Dataset ready")
+    
+    def train(self):
+        """Train the model."""
+        from transformers import (
+            TrainingArguments,
+            Trainer,
+            DataCollatorForLanguageModeling
+        )
+        
+        logger.info("Setting up training...")
+        
+        training_args = TrainingArguments(
+            output_dir=self.config.output_dir,
+            num_train_epochs=self.config.num_epochs,
+            per_device_train_batch_size=self.config.batch_size,
+            per_device_eval_batch_size=self.config.batch_size,
+            gradient_accumulation_steps=self.config.gradient_accumulation,
+            learning_rate=self.config.learning_rate,
+            weight_decay=0.01,
+            warmup_steps=self.config.warmup_steps,
+            logging_steps=10,
+            evaluation_strategy="steps",
+            eval_steps=500,
+            save_steps=500,
+            save_total_limit=3,
+            fp16=False,
+            bf16=True,
+            optim="adamw_torch",
+            gradient_checkpointing=True,
+            lr_scheduler_type="cosine",
+            report_to=["tensorboard"],
+            load_best_model_at_end=True,
+            metric_for_best_model="eval_loss",
+            greater_is_better=False,
+            push_to_hub=False
+        )
+        
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=self.tokenizer,
+            mlm=False
+        )
+        
+        trainer = Trainer(
+            model=self.model,
+            args=training_args,
+            train_dataset=self.tokenized_dataset["train"],
+            eval_dataset=self.tokenized_dataset.get("validation"),
+            tokenizer=self.tokenizer,
+            data_collator=data_collator
+        )
+        
+        logger.info("🚀 Starting training...")
+        trainer.train()
+        
+        logger.info("✅ Training complete!")
+        
+        # Save final model
+        logger.info("Saving model...")
+        trainer.save_model(self.config.output_dir)
+        self.tokenizer.save_pretrained(self.config.output_dir)
+        
+        logger.info(f"Model saved to {self.config.output_dir}")
+    
+    def push_to_hub(self):
+        """Upload model to HuggingFace Hub."""
+        from huggingface_hub import HfApi
+        
+        logger.info(f"Pushing model to {self.config.hub_model_id}...")
+        
+        api = HfApi(token=self.hf_token)
+        
+        # Create repo
+        api.create_repo(
+            self.config.hub_model_id,
+            exist_ok=True,
+            private=False
+        )
+        
+        # Upload files
+        api.upload_folder(
+            folder_path=self.config.output_dir,
+            repo_id=self.config.hub_model_id,
+            repo_type="model"
+        )
+        
+        logger.info("✅ Model pushed to Hub!")
+    
+    def run_pipeline(self) -> bool:
+        """Run complete training pipeline."""
+        try:
+            logger.info("="*60)
+            logger.info("Helion-V1.5 Training Pipeline")
+            logger.info("="*60)
+            
+            if not self.verify_setup():
+                logger.error("Setup verification failed")
+                return False
+            
+            self.prepare_model()
+            self.load_dataset()
+            self.train()
+            self.push_to_hub()
+            
+            logger.info("="*60)
+            logger.info("✅ Training pipeline completed successfully!")
+            logger.info("="*60)
+            return True
+            
+        except Exception as e:
+            logger.error(f"Training failed: {e}")
+            logger.error(traceback.format_exc())
+            return False
+
+
+def main():
+    """Main entry point."""
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Train Helion-V1.5")
+    parser.add_argument("--base-model", default="meta-llama/Llama-2-7b-hf")
+    parser.add_argument("--dataset", required=True)
+    parser.add_argument("--output-dir", default="./helion-v1.5-output")
+    parser.add_argument("--hub-model-id", default="DeepXR/Helion-V1.5")
+    parser.add_argument("--epochs", type=int, default=3)
+    parser.add_argument("--batch-size", type=int, default=4)
+    parser.add_argument("--learning-rate", type=float, default=2e-5)
+    parser.add_argument("--token", help="HuggingFace token")
+    
+    args = parser.parse_args()
+    
+    config = HelionV15Config(
+        base_model=args.base_model,
+        dataset_name=args.dataset,
+        output_dir=args.output_dir,
+        hub_model_id=args.hub_model_id,
+        num_epochs=args.epochs,
+        batch_size=args.batch_size,
+        learning_rate=args.learning_rate,
+        hf_token=args.token
+    )
+    
+    trainer = HelionV15Trainer(config)
+    success = trainer.run_pipeline()
+    
+    sys.exit(0 if success else 1)
+
+
+if __name__ == "__main__":
+    main()