isa-model 0.3.91__py3-none-any.whl → 0.4.3__py3-none-any.whl

isa_model/training/core/utils.py

@@ -1,213 +0,0 @@
-"""
-Training Utilities
-
-Helper functions and utilities for training operations.
-"""
-
-import os
-import json
-import logging
-import datetime
-from typing import Dict, Any, Optional, List
-from pathlib import Path
-
-logger = logging.getLogger(__name__)
-
-
-class TrainingUtils:
-    """Utility functions for training operations."""
-    
-    @staticmethod
-    def generate_output_dir(model_name: str, training_type: str, base_dir: str = "training_outputs") -> str:
-        """Generate a timestamped output directory."""
-        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-        safe_model_name = model_name.replace("/", "_").replace(":", "_")
-        output_dir = os.path.join(base_dir, f"{safe_model_name}_{training_type}_{timestamp}")
-        return output_dir
-    
-    @staticmethod
-    def save_training_args(args: Dict[str, Any], output_dir: str) -> None:
-        """Save training arguments to file."""
-        args_path = Path(output_dir) / "training_args.json"
-        args_path.parent.mkdir(parents=True, exist_ok=True)
-        
-        with open(args_path, 'w') as f:
-            json.dump(args, f, indent=2, default=str)
-        
-        logger.info(f"Training arguments saved to: {args_path}")
-    
-    @staticmethod
-    def load_training_args(output_dir: str) -> Dict[str, Any]:
-        """Load training arguments from file."""
-        args_path = Path(output_dir) / "training_args.json"
-        
-        if not args_path.exists():
-            raise FileNotFoundError(f"Training args not found: {args_path}")
-        
-        with open(args_path, 'r') as f:
-            args = json.load(f)
-        
-        return args
-    
-    @staticmethod
-    def get_model_info(model_name: str) -> Dict[str, Any]:
-        """Get information about a model."""
-        try:
-            from transformers import AutoConfig
-            
-            config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
-            
-            model_info = {
-                "model_name": model_name,
-                "model_type": config.model_type,
-                "vocab_size": getattr(config, 'vocab_size', None),
-                "hidden_size": getattr(config, 'hidden_size', None),
-                "num_layers": getattr(config, 'num_hidden_layers', None),
-                "num_attention_heads": getattr(config, 'num_attention_heads', None),
-                "max_position_embeddings": getattr(config, 'max_position_embeddings', None),
-            }
-            
-            return model_info
-            
-        except Exception as e:
-            logger.warning(f"Could not get model info for {model_name}: {e}")
-            return {"model_name": model_name, "error": str(e)}
-    
-    @staticmethod
-    def estimate_memory_usage(
-        model_name: str,
-        batch_size: int = 1,
-        max_length: int = 1024,
-        use_lora: bool = True
-    ) -> Dict[str, Any]:
-        """Estimate memory usage for training."""
-        try:
-            model_info = TrainingUtils.get_model_info(model_name)
-            
-            # Rough estimation based on model parameters
-            hidden_size = model_info.get('hidden_size', 4096)
-            num_layers = model_info.get('num_layers', 32)
-            vocab_size = model_info.get('vocab_size', 32000)
-            
-            # Estimate model parameters (in millions)
-            param_count = (hidden_size * hidden_size * 12 * num_layers + vocab_size * hidden_size) / 1e6
-            
-            # Base memory for model (assuming fp16)
-            model_memory_gb = param_count * 2 / 1024  # 2 bytes per parameter
-            
-            # Training memory overhead (gradients, optimizer states, activations)
-            if use_lora:
-                training_overhead = 2.0  # LoRA reduces memory usage significantly
-            else:
-                training_overhead = 4.0  # Full fine-tuning needs more memory
-            
-            # Batch and sequence length impact
-            sequence_memory = batch_size * max_length * hidden_size * 2 / (1024**3)  # Activation memory
-            
-            total_memory_gb = model_memory_gb * training_overhead + sequence_memory
-            
-            return {
-                "estimated_params_millions": param_count,
-                "model_memory_gb": model_memory_gb,
-                "total_training_memory_gb": total_memory_gb,
-                "recommended_gpu": TrainingUtils._recommend_gpu(total_memory_gb),
-                "use_lora": use_lora,
-                "batch_size": batch_size,
-                "max_length": max_length
-            }
-            
-        except Exception as e:
-            logger.warning(f"Could not estimate memory usage: {e}")
-            return {"error": str(e)}
-    
-    @staticmethod
-    def _recommend_gpu(memory_gb: float) -> str:
-        """Recommend GPU based on memory requirements."""
-        if memory_gb <= 8:
-            return "RTX 3080/4070 (8-12GB)"
-        elif memory_gb <= 16:
-            return "RTX 4080/4090 (16GB)"
-        elif memory_gb <= 24:
-            return "RTX A6000/4090 (24GB)"
-        elif memory_gb <= 40:
-            return "A100 40GB"
-        elif memory_gb <= 80:
-            return "A100 80GB"
-        else:
-            return "Multiple A100 80GB (Multi-GPU required)"
-    
-    @staticmethod
-    def validate_training_config(config: Dict[str, Any]) -> List[str]:
-        """Validate training configuration and return any issues."""
-        issues = []
-        
-        # Check required fields
-        required_fields = ["model_name", "output_dir"]
-        for field in required_fields:
-            if field not in config:
-                issues.append(f"Missing required field: {field}")
-        
-        # Check batch size
-        if config.get("batch_size", 0) <= 0:
-            issues.append("batch_size must be positive")
-        
-        # Check learning rate
-        lr = config.get("learning_rate", 0)
-        if lr <= 0 or lr > 1:
-            issues.append("learning_rate should be between 0 and 1")
-        
-        # Check epochs
-        if config.get("num_epochs", 0) <= 0:
-            issues.append("num_epochs must be positive")
-        
-        # Check LoRA config
-        if config.get("use_lora", False):
-            lora_rank = config.get("lora_rank", 8)
-            if lora_rank <= 0 or lora_rank > 256:
-                issues.append("lora_rank should be between 1 and 256")
-        
-        return issues
-    
-    @staticmethod
-    def format_training_summary(
-        config: Dict[str, Any],
-        model_info: Dict[str, Any],
-        memory_estimate: Dict[str, Any]
-    ) -> str:
-        """Format a training summary for display."""
-        summary = []
-        summary.append("=" * 60)
-        summary.append("TRAINING CONFIGURATION SUMMARY")
-        summary.append("=" * 60)
-        
-        # Model information
-        summary.append(f"Model: {config.get('model_name', 'Unknown')}")
-        summary.append(f"Model Type: {model_info.get('model_type', 'Unknown')}")
-        summary.append(f"Parameters: ~{memory_estimate.get('estimated_params_millions', 0):.1f}M")
-        
-        # Training configuration
-        summary.append(f"\nTraining Configuration:")
-        summary.append(f"  Training Type: {config.get('training_type', 'sft')}")
-        summary.append(f"  Epochs: {config.get('num_epochs', 3)}")
-        summary.append(f"  Batch Size: {config.get('batch_size', 4)}")
-        summary.append(f"  Learning Rate: {config.get('learning_rate', 2e-5)}")
-        summary.append(f"  Max Length: {config.get('max_length', 1024)}")
-        
-        # LoRA configuration
-        if config.get('use_lora', True):
-            summary.append(f"\nLoRA Configuration:")
-            summary.append(f"  LoRA Rank: {config.get('lora_rank', 8)}")
-            summary.append(f"  LoRA Alpha: {config.get('lora_alpha', 16)}")
-            summary.append(f"  LoRA Dropout: {config.get('lora_dropout', 0.05)}")
-        
-        # Memory estimation
-        summary.append(f"\nMemory Estimation:")
-        summary.append(f"  Estimated Memory: ~{memory_estimate.get('total_training_memory_gb', 0):.1f}GB")
-        summary.append(f"  Recommended GPU: {memory_estimate.get('recommended_gpu', 'Unknown')}")
-        
-        # Output
-        summary.append(f"\nOutput Directory: {config.get('output_dir', 'Unknown')}")
-        
-        summary.append("=" * 60)
-        
-        return "\n".join(summary) 

isa_model/training/factory.py

@@ -1,424 +0,0 @@
-"""
-ISA Model Training Factory
-
-A clean, simplified training factory that uses HuggingFace Transformers directly
-without external dependencies like LlamaFactory.
-"""
-
-import os
-import logging
-from typing import Optional, Dict, Any, Union, List
-from pathlib import Path
-import datetime
-
-from .core import (
-    TrainingConfig, 
-    LoRAConfig, 
-    DatasetConfig,
-    BaseTrainer,
-    SFTTrainer,
-    TrainingUtils,
-    DatasetManager,
-)
-from .cloud import TrainingJobOrchestrator
-
-logger = logging.getLogger(__name__)
-
-
-class TrainingFactory:
-    """
-    Unified Training Factory for ISA Model SDK
-    
-    Provides a clean interface for:
-    - Local training with SFT (Supervised Fine-Tuning)
-    - Cloud training on RunPod
-    - Model evaluation and management
-    
-    Example usage:
-        ```python
-        from isa_model.training import TrainingFactory
-        
-        factory = TrainingFactory()
-        
-        # Local training
-        model_path = factory.train_model(
-            model_name="google/gemma-2-4b-it",
-            dataset_path="tatsu-lab/alpaca",
-            use_lora=True,
-            num_epochs=3
-        )
-        
-        # Cloud training on RunPod
-        result = factory.train_on_runpod(
-            model_name="google/gemma-2-4b-it",
-            dataset_path="tatsu-lab/alpaca",
-            runpod_api_key="your-api-key",
-            template_id="your-template-id"
-        )
-        ```
-    """
-    
-    def __init__(self, base_output_dir: Optional[str] = None):
-        """
-        Initialize the training factory.
-        
-        Args:
-            base_output_dir: Base directory for training outputs
-        """
-        self.base_output_dir = base_output_dir or os.path.join(os.getcwd(), "training_outputs")
-        os.makedirs(self.base_output_dir, exist_ok=True)
-        
-        logger.info(f"TrainingFactory initialized with output dir: {self.base_output_dir}")
-    
-    def train_model(
-        self,
-        model_name: str,
-        dataset_path: str,
-        output_dir: Optional[str] = None,
-        training_type: str = "sft",
-        dataset_format: str = "alpaca",
-        use_lora: bool = True,
-        batch_size: int = 4,
-        num_epochs: int = 3,
-        learning_rate: float = 2e-5,
-        max_length: int = 1024,
-        lora_rank: int = 8,
-        lora_alpha: int = 16,
-        validation_split: float = 0.1,
-        **kwargs
-    ) -> str:
-        """
-        Train a model locally.
-        
-        Args:
-            model_name: Model identifier (e.g., "google/gemma-2-4b-it")
-            dataset_path: Path to dataset or HuggingFace dataset name
-            output_dir: Custom output directory
-            training_type: Type of training ("sft" supported)
-            dataset_format: Dataset format ("alpaca", "sharegpt", "custom")
-            use_lora: Whether to use LoRA for efficient training
-            batch_size: Training batch size
-            num_epochs: Number of training epochs
-            learning_rate: Learning rate
-            max_length: Maximum sequence length
-            lora_rank: LoRA rank parameter
-            lora_alpha: LoRA alpha parameter
-            validation_split: Fraction of data for validation
-            **kwargs: Additional training parameters
-            
-        Returns:
-            Path to the trained model
-            
-        Example:
-            ```python
-            model_path = factory.train_model(
-                model_name="google/gemma-2-4b-it",
-                dataset_path="tatsu-lab/alpaca",
-                use_lora=True,
-                num_epochs=3,
-                batch_size=4
-            )
-            ```
-        """
-        # Generate output directory if not provided
-        if not output_dir:
-            output_dir = TrainingUtils.generate_output_dir(
-                model_name, training_type, self.base_output_dir
-            )
-        
-        # Create configurations
-        lora_config = LoRAConfig(
-            use_lora=use_lora,
-            lora_rank=lora_rank,
-            lora_alpha=lora_alpha
-        ) if use_lora else None
-        
-        dataset_config = DatasetConfig(
-            dataset_path=dataset_path,
-            dataset_format=dataset_format,
-            max_length=max_length,
-            validation_split=validation_split
-        )
-        
-        training_config = TrainingConfig(
-            model_name=model_name,
-            output_dir=output_dir,
-            training_type=training_type,
-            num_epochs=num_epochs,
-            batch_size=batch_size,
-            learning_rate=learning_rate,
-            lora_config=lora_config,
-            dataset_config=dataset_config,
-            **kwargs
-        )
-        
-        # Print training summary
-        model_info = TrainingUtils.get_model_info(model_name)
-        memory_estimate = TrainingUtils.estimate_memory_usage(
-            model_name, batch_size, max_length, use_lora
-        )
-        
-        summary = TrainingUtils.format_training_summary(
-            training_config.to_dict(), model_info, memory_estimate
-        )
-        print(summary)
-        
-        # Validate configuration
-        issues = TrainingUtils.validate_training_config(training_config.to_dict())
-        if issues:
-            raise ValueError(f"Training configuration issues: {issues}")
-        
-        # Initialize trainer based on training type
-        if training_type.lower() == "sft":
-            trainer = SFTTrainer(training_config)
-        else:
-            raise ValueError(f"Training type '{training_type}' not supported yet")
-        
-        # Execute training
-        logger.info(f"Starting {training_type.upper()} training...")
-        result_path = trainer.train()
-        
-        logger.info(f"Training completed! Model saved to: {result_path}")
-        return result_path
-    
-    def train_on_runpod(
-        self,
-        model_name: str,
-        dataset_path: str,
-        runpod_api_key: str,
-        template_id: str,
-        gpu_type: str = "NVIDIA RTX A6000",
-        storage_config: Optional[Dict[str, Any]] = None,
-        job_name: Optional[str] = None,
-        **training_params
-    ) -> Dict[str, Any]:
-        """
-        Train a model on RunPod cloud infrastructure.
-        
-        Args:
-            model_name: Model identifier
-            dataset_path: Dataset path or HuggingFace dataset name
-            runpod_api_key: RunPod API key
-            template_id: RunPod template ID
-            gpu_type: GPU type to use
-            storage_config: Optional cloud storage configuration
-            job_name: Optional job name
-            **training_params: Additional training parameters
-            
-        Returns:
-            Training job results
-            
-        Example:
-            ```python
-            result = factory.train_on_runpod(
-                model_name="google/gemma-2-4b-it",
-                dataset_path="tatsu-lab/alpaca",
-                runpod_api_key="your-api-key",
-                template_id="your-template-id",
-                use_lora=True,
-                num_epochs=3
-            )
-            ```
-        """
-        # Import cloud components
-        from .cloud import TrainingJobOrchestrator
-        from .cloud.runpod_trainer import RunPodConfig
-        from .cloud.storage_manager import StorageConfig
-        from .cloud.job_orchestrator import JobConfig
-        
-        # Create RunPod configuration
-        runpod_config = RunPodConfig(
-            api_key=runpod_api_key,
-            template_id=template_id,
-            gpu_type=gpu_type
-        )
-        
-        # Create storage configuration if provided
-        storage_cfg = None
-        if storage_config:
-            storage_cfg = StorageConfig(**storage_config)
-        
-        # Create job configuration
-        job_config = JobConfig(
-            model_name=model_name,
-            dataset_source=dataset_path,
-            job_name=job_name or f"gemma-training-{int(datetime.datetime.now().timestamp())}",
-            **training_params
-        )
-        
-        # Initialize orchestrator and execute training
-        orchestrator = TrainingJobOrchestrator(
-            runpod_config=runpod_config,
-            storage_config=storage_cfg
-        )
-        
-        logger.info(f"Starting RunPod training for {model_name}")
-        result = orchestrator.execute_training_workflow(job_config)
-        
-        return result
-    
-    async def upload_to_huggingface(
-        self,
-        model_path: str,
-        hf_model_name: str,
-        hf_token: Optional[str] = None,
-        metadata: Optional[Dict[str, Any]] = None
-    ) -> str:
-        """
-        Upload a trained model to HuggingFace Hub using HuggingFaceStorage.
-        
-        Args:
-            model_path: Path to the trained model
-            hf_model_name: Name for the model on HuggingFace Hub
-            hf_token: HuggingFace token
-            metadata: Additional metadata for the model
-            
-        Returns:
-            URL of the uploaded model
-        """
-        try:
-            from ..core.storage.hf_storage import HuggingFaceStorage
-            
-            logger.info(f"Uploading model to HuggingFace: {hf_model_name}")
-            
-            # Initialize HuggingFace storage
-            storage = HuggingFaceStorage(
-                username="xenobordom",
-                token=hf_token
-            )
-            
-            # Prepare metadata
-            upload_metadata = metadata or {}
-            upload_metadata.update({
-                "description": f"Fine-tuned model: {hf_model_name}",
-                "training_framework": "ISA Model SDK",
-                "uploaded_from": "training_factory"
-            })
-            
-            # Upload model
-            success = await storage.save_model(
-                model_id=hf_model_name,
-                model_path=model_path,
-                metadata=upload_metadata
-            )
-            
-            if success:
-                model_url = storage.get_public_url(hf_model_name)
-                logger.info(f"Model uploaded successfully: {model_url}")
-                return model_url
-            else:
-                raise Exception("Failed to upload model")
-            
-        except Exception as e:
-            logger.error(f"Failed to upload to HuggingFace: {e}")
-            raise
-    
-    def get_training_status(self, output_dir: str) -> Dict[str, Any]:
-        """
-        Get training status from output directory.
-        
-        Args:
-            output_dir: Training output directory
-            
-        Returns:
-            Dictionary with training status information
-        """
-        status = {
-            "output_dir": output_dir,
-            "exists": os.path.exists(output_dir),
-            "files": []
-        }
-        
-        if status["exists"]:
-            status["files"] = os.listdir(output_dir)
-            
-            # Check for specific files
-            config_path = os.path.join(output_dir, "training_config.json")
-            metrics_path = os.path.join(output_dir, "training_metrics.json")
-            model_path = os.path.join(output_dir, "pytorch_model.bin")
-            
-            status["has_config"] = os.path.exists(config_path)
-            status["has_metrics"] = os.path.exists(metrics_path)
-            status["has_model"] = os.path.exists(model_path) or os.path.exists(os.path.join(output_dir, "adapter_model.bin"))
-            
-            if status["has_config"]:
-                try:
-                    status["config"] = TrainingUtils.load_training_args(output_dir)
-                except:
-                    pass
-        
-        return status
-    
-    def list_trained_models(self) -> List[Dict[str, Any]]:
-        """
-        List all trained models in the output directory.
-        
-        Returns:
-            List of model information dictionaries
-        """
-        models = []
-        
-        if os.path.exists(self.base_output_dir):
-            for item in os.listdir(self.base_output_dir):
-                item_path = os.path.join(self.base_output_dir, item)
-                if os.path.isdir(item_path):
-                    status = self.get_training_status(item_path)
-                    models.append({
-                        "name": item,
-                        "path": item_path,
-                        "created": datetime.datetime.fromtimestamp(
-                            os.path.getctime(item_path)
-                        ).isoformat(),
-                        "status": status
-                    })
-        
-        return sorted(models, key=lambda x: x["created"], reverse=True)
-
-
-# Convenience functions for quick access
-def train_gemma(
-    dataset_path: str,
-    model_size: str = "4b",
-    output_dir: Optional[str] = None,
-    **kwargs
-) -> str:
-    """
-    Quick function to train Gemma models.
-    
-    Args:
-        dataset_path: Path to training dataset
-        model_size: Model size ("2b", "4b", "7b")
-        output_dir: Output directory
-        **kwargs: Additional training parameters
-        
-    Returns:
-        Path to trained model
-        
-    Example:
-        ```python
-        from isa_model.training import train_gemma
-        
-        model_path = train_gemma(
-            dataset_path="tatsu-lab/alpaca",
-            model_size="4b",
-            num_epochs=3,
-            batch_size=4
-        )
-        ```
-    """
-    factory = TrainingFactory()
-    
-    model_map = {
-        "2b": "google/gemma-2-2b-it",
-        "4b": "google/gemma-2-4b-it", 
-        "7b": "google/gemma-2-7b-it"
-    }
-    
-    model_name = model_map.get(model_size, "google/gemma-2-4b-it")
-    
-    return factory.train_model(
-        model_name=model_name,
-        dataset_path=dataset_path,
-        output_dir=output_dir,
-        **kwargs
-    )