isa-model 0.0.1__py3-none-any.whl

isa_model/training/engine/llama_factory/examples/finetune_with_tracking.py

@@ -0,0 +1,185 @@
+"""
+Example of fine-tuning with LlamaFactory and MLflow tracking.
+
+This example demonstrates how to use LlamaFactory for fine-tuning
+and the MLflow tracking system to monitor and log the process.
+"""
+
+import os
+import argparse
+import logging
+from typing import Dict, Any
+
+from app.services.ai.models.training.engine.llama_factory import (
+    LlamaFactory,
+    DatasetFormat,
+    TrainingStrategy
+)
+from app.services.ai.models.mlops import (
+    TrainingTracker,
+    ModelStage
+)
+
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(description="Fine-tune with LlamaFactory and MLflow tracking")
+    
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        required=True,
+        help="Path or name of the base model"
+    )
+    parser.add_argument(
+        "--train_data",
+        type=str,
+        required=True,
+        help="Path to the training data"
+    )
+    parser.add_argument(
+        "--val_data",
+        type=str,
+        help="Path to the validation data"
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="outputs",
+        help="Directory to save outputs"
+    )
+    parser.add_argument(
+        "--dataset_format",
+        type=str,
+        choices=["alpaca", "sharegpt", "custom"],
+        default="alpaca",
+        help="Format of the dataset"
+    )
+    parser.add_argument(
+        "--use_lora",
+        action="store_true",
+        help="Whether to use LoRA for training"
+    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=8,
+        help="Training batch size"
+    )
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        default=3,
+        help="Number of training epochs"
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=2e-5,
+        help="Learning rate"
+    )
+    parser.add_argument(
+        "--max_length",
+        type=int,
+        default=1024,
+        help="Maximum sequence length"
+    )
+    parser.add_argument(
+        "--lora_rank",
+        type=int,
+        default=8,
+        help="LoRA rank parameter"
+    )
+    parser.add_argument(
+        "--tracking_uri",
+        type=str,
+        help="URI for MLflow tracking server"
+    )
+    parser.add_argument(
+        "--register_model",
+        action="store_true",
+        help="Whether to register the model in the registry"
+    )
+    
+    return parser.parse_args()
+
+
+def main():
+    """Run the fine-tuning process with tracking."""
+    args = parse_args()
+    
+    # Map dataset format string to enum
+    dataset_format_map = {
+        "alpaca": DatasetFormat.ALPACA,
+        "sharegpt": DatasetFormat.SHAREGPT,
+        "custom": DatasetFormat.CUSTOM
+    }
+    dataset_format = dataset_format_map[args.dataset_format]
+    
+    # Create output directory
+    os.makedirs(args.output_dir, exist_ok=True)
+    
+    # Initialize LlamaFactory
+    factory = LlamaFactory(base_output_dir=args.output_dir)
+    
+    # Initialize training tracker
+    tracker = TrainingTracker(tracking_uri=args.tracking_uri)
+    
+    # Get model name from path
+    model_name = os.path.basename(args.model_path)
+    
+    # Set up training parameters
+    train_params = {
+        "model_path": args.model_path,
+        "train_data": args.train_data,
+        "val_data": args.val_data,
+        "output_dir": None,  # Will be set by factory
+        "dataset_format": dataset_format,
+        "use_lora": args.use_lora,
+        "batch_size": args.batch_size,
+        "num_epochs": args.num_epochs,
+        "learning_rate": args.learning_rate,
+        "max_length": args.max_length,
+        "lora_rank": args.lora_rank
+    }
+    
+    # Track the training run with MLflow
+    with tracker.track_training_run(
+        model_name=model_name,
+        training_params=train_params,
+        description=f"Fine-tuning {model_name} with LlamaFactory"
+    ) as run_info:
+        # Run the fine-tuning
+        try:
+            model_path = factory.finetune(**train_params)
+            
+            # Log success
+            tracker.log_metrics({"success": 1.0})
+            logger.info(f"Fine-tuning completed successfully. Model saved to {model_path}")
+            
+            # Register the model if requested
+            if args.register_model:
+                version = tracker.register_trained_model(
+                    model_path=model_path,
+                    description=f"Fine-tuned {model_name}",
+                    stage=ModelStage.STAGING
+                )
+                logger.info(f"Model registered as version {version}")
+                
+        except Exception as e:
+            # Log failure
+            tracker.log_metrics({"success": 0.0})
+            logger.error(f"Fine-tuning failed: {e}")
+            raise
+
+
+if __name__ == "__main__":
+    main() 

isa_model/training/engine/llama_factory/examples/rlhf_with_tracking.py

@@ -0,0 +1,163 @@
+"""
+Example of RLHF (Reinforcement Learning from Human Feedback) with LlamaFactory and MLflow tracking.
+
+This example demonstrates how to use LlamaFactory for RLHF
+and the MLflow tracking system to monitor and log the process.
+"""
+
+import os
+import argparse
+import logging
+from typing import Dict, Any
+
+from app.services.ai.models.training.engine.llama_factory import (
+    LlamaFactory,
+    TrainingStrategy
+)
+from app.services.ai.models.mlops import (
+    TrainingTracker,
+    ModelStage
+)
+
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(description="Run RLHF with LlamaFactory and MLflow tracking")
+    
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        required=True,
+        help="Path or name of the base model"
+    )
+    parser.add_argument(
+        "--reward_model",
+        type=str,
+        required=True,
+        help="Path to the reward model"
+    )
+    parser.add_argument(
+        "--train_data",
+        type=str,
+        required=True,
+        help="Path to the training data"
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="outputs",
+        help="Directory to save outputs"
+    )
+    parser.add_argument(
+        "--use_lora",
+        action="store_true",
+        help="Whether to use LoRA for training"
+    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=4,
+        help="Training batch size"
+    )
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        default=1,
+        help="Number of training epochs"
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=1e-5,
+        help="Learning rate"
+    )
+    parser.add_argument(
+        "--kl_coef",
+        type=float,
+        default=0.1,
+        help="KL coefficient for PPO"
+    )
+    parser.add_argument(
+        "--tracking_uri",
+        type=str,
+        help="URI for MLflow tracking server"
+    )
+    parser.add_argument(
+        "--register_model",
+        action="store_true",
+        help="Whether to register the model in the registry"
+    )
+    
+    return parser.parse_args()
+
+
+def main():
+    """Run the RLHF process with tracking."""
+    args = parse_args()
+    
+    # Create output directory
+    os.makedirs(args.output_dir, exist_ok=True)
+    
+    # Initialize LlamaFactory
+    factory = LlamaFactory(base_output_dir=args.output_dir)
+    
+    # Initialize training tracker
+    tracker = TrainingTracker(tracking_uri=args.tracking_uri)
+    
+    # Get model name from path
+    model_name = os.path.basename(args.model_path)
+    
+    # Set up RLHF parameters
+    rl_params = {
+        "model_path": args.model_path,
+        "reward_model": args.reward_model,
+        "train_data": args.train_data,
+        "output_dir": None,  # Will be set by factory
+        "use_lora": args.use_lora,
+        "batch_size": args.batch_size,
+        "num_epochs": args.num_epochs,
+        "learning_rate": args.learning_rate,
+        "kl_coef": args.kl_coef
+    }
+    
+    # Track the RLHF run with MLflow
+    with tracker.track_training_run(
+        model_name=model_name,
+        training_params=rl_params,
+        description=f"RLHF for {model_name} with LlamaFactory",
+        experiment_type="rl"
+    ) as run_info:
+        # Run the RLHF
+        try:
+            model_path = factory.rlhf(**rl_params)
+            
+            # Log success
+            tracker.log_metrics({"success": 1.0})
+            logger.info(f"RLHF completed successfully. Model saved to {model_path}")
+            
+            # Register the model if requested
+            if args.register_model:
+                version = tracker.register_trained_model(
+                    model_path=model_path,
+                    description=f"RLHF-tuned {model_name}",
+                    stage=ModelStage.STAGING
+                )
+                logger.info(f"Model registered as version {version}")
+                
+        except Exception as e:
+            # Log failure
+            tracker.log_metrics({"success": 0.0})
+            logger.error(f"RLHF failed: {e}")
+            raise
+
+
+if __name__ == "__main__":
+    main() 

isa_model/training/engine/llama_factory/factory.py

@@ -0,0 +1,331 @@
+"""
+LlamaFactory interface module.
+
+This module provides the main interface for using LlamaFactory functionality.
+"""
+
+import os
+import logging
+from typing import Optional, Dict, Any, Union, List
+
+from .config import (
+    LlamaFactoryConfig,
+    SFTConfig,
+    RLConfig,
+    DPOConfig,
+    TrainingStrategy,
+    DatasetFormat,
+    create_default_config
+)
+from .trainer import LlamaFactoryTrainer
+from .rl import LlamaFactoryRL
+from .data_adapter import DataAdapterFactory
+
+
+logger = logging.getLogger(__name__)
+
+
+class LlamaFactory:
+    """
+    Main interface class for LlamaFactory operations.
+    
+    This class provides a simplified interface to the LlamaFactory functionality
+    for fine-tuning and reinforcement learning.
+    
+    Example for fine-tuning:
+        ```python
+        # Create a LlamaFactory instance
+        factory = LlamaFactory()
+        
+        # Fine-tune a model
+        model_path = factory.finetune(
+            model_path="meta-llama/Llama-2-7b-hf",
+            train_data="path/to/data.json",
+            val_data="path/to/val_data.json",  # Optional
+            output_dir="path/to/output",
+            dataset_format=DatasetFormat.ALPACA,
+            use_lora=True,
+            num_epochs=3
+        )
+        ```
+        
+    Example for RL training:
+        ```python
+        # Create a LlamaFactory instance
+        factory = LlamaFactory()
+        
+        # Train with DPO
+        model_path = factory.dpo(
+            model_path="meta-llama/Llama-2-7b-hf",
+            train_data="path/to/preferences.json",
+            output_dir="path/to/output",
+            reference_model="meta-llama/Llama-2-7b-hf",  # Optional
+            beta=0.1
+        )
+        ```
+    """
+    
+    def __init__(self, base_output_dir: Optional[str] = None):
+        """
+        Initialize the LlamaFactory interface.
+        
+        Args:
+            base_output_dir: Base directory for 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)
+        
+    def _get_output_dir(self, name: str, output_dir: Optional[str] = None) -> str:
+        """
+        Get the output directory for training.
+        
+        Args:
+            name: Name for the output directory
+            output_dir: Optional specific output directory
+            
+        Returns:
+            Output directory path
+        """
+        if output_dir:
+            return output_dir
+            
+        import datetime
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        return os.path.join(self.base_output_dir, f"{name}_{timestamp}")
+    
+    def finetune(
+        self,
+        model_path: str,
+        train_data: str,
+        val_data: Optional[str] = None,
+        output_dir: Optional[str] = None,
+        dataset_format: DatasetFormat = DatasetFormat.ALPACA,
+        use_lora: bool = True,
+        batch_size: int = 8,
+        num_epochs: int = 3,
+        learning_rate: float = 2e-5,
+        lora_rank: int = 8,
+        lora_alpha: int = 16,
+        max_length: int = 1024,
+        **kwargs
+    ) -> str:
+        """
+        Fine-tune a model using supervised learning.
+        
+        Args:
+            model_path: Path or name of the base model
+            train_data: Path to the training data
+            val_data: Path to the validation data (optional)
+            output_dir: Directory to save outputs
+            dataset_format: Format of the dataset
+            use_lora: Whether to use LoRA for training
+            batch_size: Training batch size
+            num_epochs: Number of training epochs
+            learning_rate: Learning rate
+            lora_rank: LoRA rank parameter
+            lora_alpha: LoRA alpha parameter
+            max_length: Maximum sequence length
+            **kwargs: Additional parameters for SFTConfig
+            
+        Returns:
+            Path to the trained model
+        """
+        # Check if data conversion is needed and convert
+        adapter = DataAdapterFactory.create_adapter(dataset_format)
+        converted_train_data = adapter.convert_data(train_data)
+        converted_val_data = adapter.convert_data(val_data) if val_data else None
+        
+        # Create configuration
+        output_dir = self._get_output_dir("sft", output_dir)
+        config = SFTConfig(
+            model_path=model_path,
+            train_file=converted_train_data,
+            val_file=converted_val_data,
+            output_dir=output_dir,
+            use_lora=use_lora,
+            batch_size=batch_size,
+            num_epochs=num_epochs,
+            learning_rate=learning_rate,
+            lora_rank=lora_rank,
+            lora_alpha=lora_alpha,
+            max_length=max_length,
+            dataset_format=dataset_format,
+            **kwargs
+        )
+        
+        # Initialize and run trainer
+        trainer = LlamaFactoryTrainer(config)
+        model_dir = trainer.train()
+        
+        # Export model if using LoRA
+        if use_lora:
+            model_dir = trainer.export_model()
+            
+        return model_dir
+    
+    def rlhf(
+        self,
+        model_path: str,
+        reward_model: str,
+        train_data: str,
+        output_dir: Optional[str] = None,
+        use_lora: bool = True,
+        batch_size: int = 4,
+        num_epochs: int = 1,
+        learning_rate: float = 1e-5,
+        kl_coef: float = 0.1,
+        **kwargs
+    ) -> str:
+        """
+        Train a model with RLHF (Reinforcement Learning from Human Feedback).
+        
+        Args:
+            model_path: Path or name of the base model
+            reward_model: Path to the reward model
+            train_data: Path to the training data
+            output_dir: Directory to save outputs
+            use_lora: Whether to use LoRA for training
+            batch_size: Training batch size
+            num_epochs: Number of training epochs
+            learning_rate: Learning rate
+            kl_coef: KL coefficient for PPO
+            **kwargs: Additional parameters for RLConfig
+            
+        Returns:
+            Path to the trained model
+        """
+        # Create configuration
+        output_dir = self._get_output_dir("rlhf", output_dir)
+        config = RLConfig(
+            model_path=model_path,
+            reward_model=reward_model,
+            train_file=train_data,
+            output_dir=output_dir,
+            use_lora=use_lora,
+            batch_size=batch_size,
+            num_epochs=num_epochs,
+            learning_rate=learning_rate,
+            kl_coef=kl_coef,
+            **kwargs
+        )
+        
+        # Initialize and run RL trainer
+        rl_trainer = LlamaFactoryRL(config)
+        model_dir = rl_trainer.train()
+        
+        # Export model if using LoRA
+        if use_lora:
+            model_dir = rl_trainer.export_model()
+            
+        return model_dir
+    
+    def dpo(
+        self,
+        model_path: str,
+        train_data: str,
+        val_data: Optional[str] = None,
+        reference_model: Optional[str] = None,
+        output_dir: Optional[str] = None,
+        use_lora: bool = True,
+        batch_size: int = 4,
+        num_epochs: int = 3,
+        learning_rate: float = 5e-6,
+        beta: float = 0.1,
+        **kwargs
+    ) -> str:
+        """
+        Train a model with DPO (Direct Preference Optimization).
+        
+        Args:
+            model_path: Path or name of the base model
+            train_data: Path to the training data
+            val_data: Path to the validation data (optional)
+            reference_model: Path to the reference model (optional)
+            output_dir: Directory to save outputs
+            use_lora: Whether to use LoRA for training
+            batch_size: Training batch size
+            num_epochs: Number of training epochs
+            learning_rate: Learning rate
+            beta: DPO beta parameter
+            **kwargs: Additional parameters for DPOConfig
+            
+        Returns:
+            Path to the trained model
+        """
+        # Create configuration
+        output_dir = self._get_output_dir("dpo", output_dir)
+        config = DPOConfig(
+            model_path=model_path,
+            train_file=train_data,
+            val_file=val_data,
+            reference_model=reference_model,
+            output_dir=output_dir,
+            use_lora=use_lora,
+            batch_size=batch_size,
+            num_epochs=num_epochs,
+            learning_rate=learning_rate,
+            beta=beta,
+            **kwargs
+        )
+        
+        # Initialize and run DPO trainer
+        dpo_trainer = LlamaFactoryRL(config)
+        model_dir = dpo_trainer.train()
+        
+        # Export model if using LoRA
+        if use_lora:
+            model_dir = dpo_trainer.export_model()
+            
+        return model_dir
+    
+    def train_reward_model(
+        self,
+        model_path: str,
+        train_data: str,
+        val_data: Optional[str] = None,
+        output_dir: Optional[str] = None,
+        use_lora: bool = True,
+        batch_size: int = 8,
+        num_epochs: int = 3,
+        learning_rate: float = 1e-5,
+        **kwargs
+    ) -> str:
+        """
+        Train a reward model for RLHF.
+        
+        Args:
+            model_path: Path or name of the base model
+            train_data: Path to the training data with preferences
+            val_data: Path to the validation data (optional)
+            output_dir: Directory to save outputs
+            use_lora: Whether to use LoRA for training
+            batch_size: Training batch size
+            num_epochs: Number of training epochs
+            learning_rate: Learning rate
+            **kwargs: Additional parameters for RLConfig
+            
+        Returns:
+            Path to the trained reward model
+        """
+        # Create temporary RL config
+        output_dir = self._get_output_dir("reward_model", output_dir)
+        config = RLConfig(
+            model_path=model_path,
+            train_file=train_data,
+            output_dir=output_dir,
+            use_lora=use_lora,
+            batch_size=batch_size,
+            num_epochs=num_epochs,
+            learning_rate=learning_rate,
+            **kwargs
+        )
+        
+        # Initialize RL trainer and train reward model
+        rl_trainer = LlamaFactoryRL(config)
+        model_dir = rl_trainer.train_reward_model()
+        
+        # Export model if using LoRA
+        if use_lora:
+            model_dir = rl_trainer.export_model()
+            
+        return model_dir