isa-model 0.0.2__py3-none-any.whl → 0.3.1__py3-none-any.whl

isa_model/training/factory.py

@@ -1,10 +1,8 @@
 """
-Unified Training Factory for ISA Model Framework
+ISA Model Training Factory
 
-This factory provides a single interface for all training operations:
-- LLM fine-tuning (SFT, DPO, RLHF)
-- Image model training (Flux, LoRA)
-- Model evaluation and benchmarking
+A clean, simplified training factory that uses HuggingFace Transformers directly
+without external dependencies like LlamaFactory.
 """
 
 import os
@@ -13,43 +11,49 @@ from typing import Optional, Dict, Any, Union, List
 from pathlib import Path
 import datetime
 
-from .engine.llama_factory import LlamaFactory, TrainingStrategy, DatasetFormat
-from .engine.llama_factory.config import SFTConfig, RLConfig, DPOConfig
+from .core import (
+    TrainingConfig, 
+    LoRAConfig, 
+    DatasetConfig,
+    BaseTrainer,
+    SFTTrainer,
+    TrainingUtils,
+    DatasetManager,
+)
+from .cloud import TrainingJobOrchestrator
 
 logger = logging.getLogger(__name__)
 
 
 class TrainingFactory:
     """
-    Unified factory for all AI model training operations.
+    Unified Training Factory for ISA Model SDK
     
-    This class provides simplified interfaces for:
-    - LLM training using LlamaFactory
-    - Image model training using Flux/LoRA
-    - Model evaluation and benchmarking
+    Provides a clean interface for:
+    - Local training with SFT (Supervised Fine-Tuning)
+    - Cloud training on RunPod
+    - Model evaluation and management
     
-    Example usage for fine-tuning Gemma 3:4B:
+    Example usage:
         ```python
         from isa_model.training import TrainingFactory
         
         factory = TrainingFactory()
         
-        # Fine-tune with your dataset
-        model_path = factory.finetune_llm(
+        # Local training
+        model_path = factory.train_model(
             model_name="google/gemma-2-4b-it",
-            dataset_path="path/to/your/data.json",
-            training_type="sft",
+            dataset_path="tatsu-lab/alpaca",
             use_lora=True,
-            num_epochs=3,
-            batch_size=4,
-            learning_rate=2e-5
+            num_epochs=3
         )
         
-        # Train with DPO for preference optimization
-        dpo_model = factory.train_with_preferences(
-            model_path=model_path,
-            preference_data="path/to/preferences.json",
-            beta=0.1
+        # 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"
         )
         ```
     """
@@ -59,32 +63,19 @@ class TrainingFactory:
         Initialize the training factory.
         
         Args:
-            base_output_dir: Base directory for all training outputs
+            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)
         
-        # Initialize sub-factories
-        self.llm_factory = LlamaFactory(base_output_dir=os.path.join(self.base_output_dir, "llm"))
-        
         logger.info(f"TrainingFactory initialized with output dir: {self.base_output_dir}")
     
-    def _get_output_dir(self, model_name: str, training_type: str) -> str:
-        """Generate timestamped output directory."""
-        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-        safe_model_name = model_name.replace("/", "_").replace(":", "_")
-        return os.path.join(self.base_output_dir, f"{safe_model_name}_{training_type}_{timestamp}")
-    
-    # =================
-    # LLM Training Methods
-    # =================
-    
-    def finetune_llm(
+    def train_model(
         self,
         model_name: str,
         dataset_path: str,
-        training_type: str = "sft",
         output_dir: Optional[str] = None,
+        training_type: str = "sft",
         dataset_format: str = "alpaca",
         use_lora: bool = True,
         batch_size: int = 4,
@@ -93,17 +84,17 @@ class TrainingFactory:
         max_length: int = 1024,
         lora_rank: int = 8,
         lora_alpha: int = 16,
-        val_dataset_path: Optional[str] = None,
+        validation_split: float = 0.1,
         **kwargs
     ) -> str:
         """
-        Fine-tune an LLM model.
+        Train a model locally.
         
         Args:
-            model_name: Model identifier (e.g., "google/gemma-2-4b-it", "meta-llama/Llama-2-7b-hf")
-            dataset_path: Path to training dataset
-            training_type: Type of training ("sft", "dpo", "rlhf")
+            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
@@ -112,7 +103,7 @@ class TrainingFactory:
             max_length: Maximum sequence length
             lora_rank: LoRA rank parameter
             lora_alpha: LoRA alpha parameter
-            val_dataset_path: Path to validation dataset (optional)
+            validation_split: Fraction of data for validation
             **kwargs: Additional training parameters
             
         Returns:
@@ -120,184 +111,207 @@ class TrainingFactory:
             
         Example:
             ```python
-            # Fine-tune Gemma 3:4B with your dataset
-            model_path = factory.finetune_llm(
+            model_path = factory.train_model(
                 model_name="google/gemma-2-4b-it",
-                dataset_path="my_training_data.json",
-                training_type="sft",
+                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 = self._get_output_dir(model_name, training_type)
+            output_dir = TrainingUtils.generate_output_dir(
+                model_name, training_type, self.base_output_dir
+            )
         
-        # Convert format string to enum
-        format_map = {
-            "alpaca": DatasetFormat.ALPACA,
-            "sharegpt": DatasetFormat.SHAREGPT,
-            "custom": DatasetFormat.CUSTOM
-        }
-        dataset_format_enum = format_map.get(dataset_format, DatasetFormat.ALPACA)
+        # Create configurations
+        lora_config = LoRAConfig(
+            use_lora=use_lora,
+            lora_rank=lora_rank,
+            lora_alpha=lora_alpha
+        ) if use_lora else None
         
-        if training_type.lower() == "sft":
-            return self.llm_factory.finetune(
-                model_path=model_name,
-                train_data=dataset_path,
-                val_data=val_dataset_path,
-                output_dir=output_dir,
-                dataset_format=dataset_format_enum,
-                use_lora=use_lora,
-                batch_size=batch_size,
-                num_epochs=num_epochs,
-                learning_rate=learning_rate,
-                max_length=max_length,
-                lora_rank=lora_rank,
-                lora_alpha=lora_alpha,
-                **kwargs
-            )
-        else:
-            raise ValueError(f"Training type '{training_type}' not supported yet. Use 'sft' for now.")
-    
-    def train_with_preferences(
-        self,
-        model_path: str,
-        preference_data: str,
-        output_dir: Optional[str] = None,
-        reference_model: Optional[str] = None,
-        beta: float = 0.1,
-        use_lora: bool = True,
-        batch_size: int = 4,
-        num_epochs: int = 3,
-        learning_rate: float = 5e-6,
-        val_data: Optional[str] = None,
-        **kwargs
-    ) -> str:
-        """
-        Train model with preference data using DPO.
+        dataset_config = DatasetConfig(
+            dataset_path=dataset_path,
+            dataset_format=dataset_format,
+            max_length=max_length,
+            validation_split=validation_split
+        )
         
-        Args:
-            model_path: Path to the base model
-            preference_data: Path to preference dataset
-            output_dir: Custom output directory
-            reference_model: Reference model for DPO (optional)
-            beta: DPO beta parameter
-            use_lora: Whether to use LoRA
-            batch_size: Training batch size
-            num_epochs: Number of epochs
-            learning_rate: Learning rate
-            val_data: Validation data path
-            **kwargs: Additional parameters
-            
-        Returns:
-            Path to the trained model
-        """
-        if not output_dir:
-            model_name = os.path.basename(model_path)
-            output_dir = self._get_output_dir(model_name, "dpo")
-        
-        return self.llm_factory.dpo(
-            model_path=model_path,
-            train_data=preference_data,
-            val_data=val_data,
-            reference_model=reference_model,
+        training_config = TrainingConfig(
+            model_name=model_name,
             output_dir=output_dir,
-            use_lora=use_lora,
-            batch_size=batch_size,
+            training_type=training_type,
             num_epochs=num_epochs,
+            batch_size=batch_size,
             learning_rate=learning_rate,
-            beta=beta,
+            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_reward_model(
+    def train_on_runpod(
         self,
-        model_path: str,
-        reward_data: str,
-        output_dir: Optional[str] = None,
-        use_lora: bool = True,
-        batch_size: int = 8,
-        num_epochs: int = 3,
-        learning_rate: float = 1e-5,
-        val_data: Optional[str] = None,
-        **kwargs
-    ) -> str:
+        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 reward model for RLHF.
+        Train a model on RunPod cloud infrastructure.
         
         Args:
-            model_path: Base model path
-            reward_data: Reward training data
-            output_dir: Output directory
-            use_lora: Whether to use LoRA
-            batch_size: Batch size
-            num_epochs: Number of epochs
-            learning_rate: Learning rate
-            val_data: Validation data
-            **kwargs: Additional parameters
+            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:
-            Path to trained reward model
+            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
+            )
+            ```
         """
-        if not output_dir:
-            model_name = os.path.basename(model_path)
-            output_dir = self._get_output_dir(model_name, "reward")
+        # 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
         
-        return self.llm_factory.train_reward_model(
-            model_path=model_path,
-            train_data=reward_data,
-            val_data=val_data,
-            output_dir=output_dir,
-            use_lora=use_lora,
-            batch_size=batch_size,
-            num_epochs=num_epochs,
-            learning_rate=learning_rate,
-            **kwargs
+        # 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
     
-    # =================
-    # Image Model Training Methods
-    # =================
-    
-    def train_image_model(
+    async def upload_to_huggingface(
         self,
-        model_type: str = "flux",
-        training_images_dir: str = "",
-        output_dir: Optional[str] = None,
-        use_lora: bool = True,
-        num_epochs: int = 1000,
-        batch_size: int = 1,
-        learning_rate: float = 1e-4,
-        **kwargs
+        model_path: str,
+        hf_model_name: str,
+        hf_token: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None
     ) -> str:
         """
-        Train an image generation model.
+        Upload a trained model to HuggingFace Hub using HuggingFaceStorage.
         
         Args:
-            model_type: Type of model ("flux", "lora")
-            training_images_dir: Directory containing training images
-            output_dir: Output directory
-            use_lora: Whether to use LoRA
-            num_epochs: Training epochs
-            batch_size: Batch size
-            learning_rate: Learning rate
-            **kwargs: Additional parameters
+            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:
-            Path to trained model
+            URL of the uploaded model
         """
-        if not output_dir:
-            output_dir = self._get_output_dir("image_model", model_type)
-        
-        # TODO: Implement image model training
-        logger.warning("Image model training not fully implemented yet")
-        return output_dir
-    
-    # =================
-    # Utility Methods
-    # =================
+        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]:
         """
@@ -318,6 +332,21 @@ class TrainingFactory:
         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]]:
@@ -333,26 +362,28 @@ class TrainingFactory:
             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()
+                        ).isoformat(),
+                        "status": status
                     })
         
         return sorted(models, key=lambda x: x["created"], reverse=True)
 
 
 # Convenience functions for quick access
-def finetune_gemma(
+def train_gemma(
     dataset_path: str,
     model_size: str = "4b",
     output_dir: Optional[str] = None,
     **kwargs
 ) -> str:
     """
-    Quick function to fine-tune Gemma models.
+    Quick function to train Gemma models.
     
     Args:
         dataset_path: Path to training dataset
@@ -361,14 +392,14 @@ def finetune_gemma(
         **kwargs: Additional training parameters
         
     Returns:
-        Path to fine-tuned model
+        Path to trained model
         
     Example:
         ```python
-        from isa_model.training import finetune_gemma
+        from isa_model.training import train_gemma
         
-        model_path = finetune_gemma(
-            dataset_path="my_data.json",
+        model_path = train_gemma(
+            dataset_path="tatsu-lab/alpaca",
             model_size="4b",
             num_epochs=3,
             batch_size=4
@@ -385,7 +416,7 @@ def finetune_gemma(
     
     model_name = model_map.get(model_size, "google/gemma-2-4b-it")
     
-    return factory.finetune_llm(
+    return factory.train_model(
         model_name=model_name,
         dataset_path=dataset_path,
         output_dir=output_dir,