isa-model 0.1.0__py3-none-any.whl

isa_model/scripts/training_tracker.py

@@ -0,0 +1,257 @@
+"""
+MLflow tracker for training workflows.
+"""
+
+import os
+import json
+import logging
+from typing import Dict, List, Optional, Any, Union
+from contextlib import contextmanager
+
+from .mlflow_manager import MLflowManager, ExperimentType
+from .model_registry import ModelRegistry, ModelStage
+
+
+logger = logging.getLogger(__name__)
+
+
+class TrainingTracker:
+    """
+    Tracker for model training workflows.
+    
+    This class provides utilities to track model training using MLflow
+    and register trained models in the model registry.
+    
+    Example:
+        ```python
+        # Initialize tracker
+        tracker = TrainingTracker(
+            tracking_uri="http://localhost:5000",
+            registry_uri="http://localhost:5000"
+        )
+        
+        # Start tracking training
+        with tracker.track_training_run(
+            model_name="llama-7b",
+            training_params={
+                "learning_rate": 2e-5,
+                "batch_size": 8,
+                "epochs": 3
+            }
+        ) as run_info:
+            # Train the model...
+            
+            # Log metrics during training
+            tracker.log_metrics({
+                "train_loss": 0.1,
+                "val_loss": 0.2
+            })
+            
+            # After training completes
+            model_path = "/path/to/trained_model"
+            
+            # Register the model
+            tracker.register_trained_model(
+                model_path=model_path,
+                metrics={
+                    "accuracy": 0.95,
+                    "f1": 0.92
+                },
+                stage=ModelStage.STAGING
+            )
+        ```
+    """
+    
+    def __init__(
+        self,
+        tracking_uri: Optional[str] = None,
+        artifact_uri: Optional[str] = None,
+        registry_uri: Optional[str] = None
+    ):
+        """
+        Initialize the training tracker.
+        
+        Args:
+            tracking_uri: URI for MLflow tracking server
+            artifact_uri: URI for MLflow artifacts
+            registry_uri: URI for MLflow model registry
+        """
+        self.mlflow_manager = MLflowManager(
+            tracking_uri=tracking_uri,
+            artifact_uri=artifact_uri,
+            registry_uri=registry_uri
+        )
+        self.model_registry = ModelRegistry(
+            tracking_uri=tracking_uri,
+            registry_uri=registry_uri
+        )
+        self.current_run_info = {}
+        
+    @contextmanager
+    def track_training_run(
+        self,
+        model_name: str,
+        training_params: Dict[str, Any],
+        description: Optional[str] = None,
+        tags: Optional[Dict[str, str]] = None,
+        experiment_type: ExperimentType = ExperimentType.TRAINING
+    ):
+        """
+        Track a training run with MLflow.
+        
+        Args:
+            model_name: Name of the model being trained
+            training_params: Parameters for the training run
+            description: Description of the training run
+            tags: Tags for the training run
+            experiment_type: Type of experiment
+            
+        Yields:
+            Dictionary with run information
+        """
+        run_info = {
+            "model_name": model_name,
+            "params": training_params,
+            "metrics": {}
+        }
+        
+        # Add description to tags if provided
+        if tags is None:
+            tags = {}
+            
+        if description:
+            tags["description"] = description
+            
+        # Start the MLflow run
+        with self.mlflow_manager.start_run(
+            experiment_type=experiment_type,
+            model_name=model_name,
+            tags=tags
+        ) as run:
+            run_info["run_id"] = run.info.run_id
+            run_info["experiment_id"] = run.info.experiment_id
+            run_info["status"] = "running"
+            
+            # Save parameters
+            self.mlflow_manager.log_params(training_params)
+            
+            self.current_run_info = run_info
+            try:
+                yield run_info
+                # Mark as successful if no exceptions
+                run_info["status"] = "completed"
+            except Exception as e:
+                # Mark as failed if exception occurred
+                run_info["status"] = "failed"
+                run_info["error"] = str(e)
+                self.mlflow_manager.set_tracking_tag("error", str(e))
+                raise
+            finally:
+                self.current_run_info = {}
+                
+    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
+        """
+        Log metrics to the current run.
+        
+        Args:
+            metrics: Dictionary of metrics to log
+            step: Step value for the metrics
+        """
+        self.mlflow_manager.log_metrics(metrics, step)
+        
+        if self.current_run_info:
+            if "metrics" not in self.current_run_info:
+                self.current_run_info["metrics"] = {}
+                
+            # Only keep the latest metrics
+            self.current_run_info["metrics"].update(metrics)
+            
+    def log_artifacts(self, local_dir: str, artifact_path: Optional[str] = None) -> None:
+        """
+        Log artifacts to the current run.
+        
+        Args:
+            local_dir: Local directory containing artifacts
+            artifact_path: Path for the artifacts in MLflow
+        """
+        self.mlflow_manager.log_artifacts(local_dir, artifact_path)
+        
+    def register_trained_model(
+        self,
+        model_path: str,
+        metrics: Optional[Dict[str, float]] = None,
+        description: Optional[str] = None,
+        tags: Optional[Dict[str, str]] = None,
+        stage: Optional[ModelStage] = None,
+        flavor: str = "pyfunc"
+    ) -> Optional[str]:
+        """
+        Register a trained model with MLflow.
+        
+        Args:
+            model_path: Path to the trained model
+            metrics: Evaluation metrics for the model
+            description: Description of the model
+            tags: Tags for the model
+            stage: Stage to register the model in
+            flavor: MLflow model flavor
+            
+        Returns:
+            Version of the registered model or None if registration failed
+        """
+        if not self.current_run_info:
+            logger.warning("No active run. Model cannot be registered.")
+            return None
+            
+        model_name = self.current_run_info.get("model_name")
+        if not model_name:
+            logger.warning("Model name not available in run info. Using generic name.")
+            model_name = "unnamed_model"
+            
+        # Log final metrics if provided
+        if metrics:
+            self.log_metrics(metrics)
+            
+        # Prepare model tags
+        if tags is None:
+            tags = {}
+            
+        # Add run ID to tags
+        tags["run_id"] = self.current_run_info.get("run_id", "")
+        
+        # Add metrics to tags
+        for k, v in self.current_run_info.get("metrics", {}).items():
+            tags[f"metric.{k}"] = str(v)
+            
+        # Log model to MLflow
+        artifact_path = self.mlflow_manager.log_model(
+            model_path=model_path,
+            name=model_name,
+            flavor=flavor
+        )
+        
+        if not artifact_path:
+            logger.error("Failed to log model to MLflow.")
+            return None
+            
+        # Get model URI
+        run_id = self.current_run_info.get("run_id")
+        model_uri = f"runs:/{run_id}/{artifact_path}"
+        
+        # Register the model
+        version = self.model_registry.register_model(
+            name=model_name,
+            source=model_uri,
+            description=description,
+            tags=tags
+        )
+        
+        # Transition to the specified stage if provided
+        if version and stage:
+            self.model_registry.transition_model_version_stage(
+                name=model_name,
+                version=version,
+                stage=stage
+            )
+            
+        return version 

isa_model/training/engine/llama_factory/__init__.py

@@ -0,0 +1,39 @@
+"""
+LlamaFactory engine for fine-tuning and reinforcement learning.
+
+This package provides interfaces for using LlamaFactory to:
+- Fine-tune models with various datasets
+- Perform reinforcement learning from human feedback (RLHF)
+- Support instruction tuning and preference optimization
+"""
+
+from .config import (
+    LlamaFactoryConfig,
+    SFTConfig,
+    RLConfig,
+    DPOConfig,
+    TrainingStrategy,
+    DatasetFormat,
+    create_default_config
+)
+from .trainer import LlamaFactoryTrainer
+from .rl import LlamaFactoryRL
+from .factory import LlamaFactory
+from .data_adapter import DataAdapter, AlpacaAdapter, ShareGPTAdapter, DataAdapterFactory
+
+__all__ = [
+    "LlamaFactoryTrainer",
+    "LlamaFactoryRL",
+    "LlamaFactoryConfig",
+    "SFTConfig",
+    "RLConfig",
+    "DPOConfig",
+    "TrainingStrategy",
+    "DatasetFormat",
+    "create_default_config",
+    "LlamaFactory",
+    "DataAdapter",
+    "AlpacaAdapter",
+    "ShareGPTAdapter",
+    "DataAdapterFactory"
+] 

isa_model/training/engine/llama_factory/config.py

@@ -0,0 +1,115 @@
+"""
+Configuration classes for LlamaFactory training and RL.
+"""
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
+from enum import Enum
+import os
+from pathlib import Path
+
+
+class TrainingStrategy(str, Enum):
+    """Training strategies supported by LlamaFactory."""
+    
+    SUPERVISED_FINETUNING = "sft"
+    REINFORCEMENT_LEARNING = "rl"
+    PREFERENCE_OPTIMIZATION = "dpo"
+    PREFERENCE_PAIRWISE = "ppo"
+
+
+class DatasetFormat(str, Enum):
+    """Dataset formats supported by LlamaFactory."""
+    
+    ALPACA = "alpaca"
+    SHAREGPT = "sharegpt"
+    CUSTOM = "custom"
+
+
+@dataclass
+class LlamaFactoryConfig:
+    """Base configuration for LlamaFactory trainers."""
+    
+    model_path: str
+    output_dir: str = field(default_factory=lambda: os.path.join(os.getcwd(), "outputs"))
+    
+    # Training parameters
+    batch_size: int = 8
+    num_epochs: int = 3
+    learning_rate: float = 2e-5
+    weight_decay: float = 0.01
+    lr_scheduler_type: str = "cosine"
+    warmup_ratio: float = 0.1
+    
+    # Model parameters
+    lora_rank: int = 8
+    lora_alpha: int = 16
+    lora_dropout: float = 0.05
+    use_lora: bool = True
+    
+    # Data processing
+    max_length: int = 1024
+    dataset_format: DatasetFormat = DatasetFormat.ALPACA
+    
+    # Logging
+    log_with: str = "mlflow"
+    logging_steps: int = 10
+    
+    def to_dict(self) -> Dict:
+        """Convert config to dictionary for LlamaFactory CLI."""
+        return {k: v.value if isinstance(v, Enum) else v 
+                for k, v in self.__dict__.items()}
+
+
+@dataclass
+class SFTConfig(LlamaFactoryConfig):
+    """Configuration for supervised fine-tuning."""
+    
+    strategy: TrainingStrategy = TrainingStrategy.SUPERVISED_FINETUNING
+    train_file: str = ""
+    val_file: Optional[str] = None
+    
+    # SFT specific settings
+    cutoff_len: int = 1024
+    normalize_data: bool = True
+    
+
+@dataclass
+class RLConfig(LlamaFactoryConfig):
+    """Configuration for reinforcement learning."""
+    
+    strategy: TrainingStrategy = TrainingStrategy.REINFORCEMENT_LEARNING
+    reward_model: str = ""
+    train_file: str = ""
+    
+    # RL specific settings
+    kl_coef: float = 0.1
+    top_k: int = 0
+    top_p: float = 1.0
+    temperature: float = 1.0
+    
+
+@dataclass
+class DPOConfig(LlamaFactoryConfig):
+    """Configuration for direct preference optimization."""
+    
+    strategy: TrainingStrategy = TrainingStrategy.PREFERENCE_OPTIMIZATION
+    train_file: str = ""
+    val_file: Optional[str] = None
+    
+    # DPO specific settings
+    beta: float = 0.1
+    reference_model: Optional[str] = None
+    
+
+def create_default_config(strategy: TrainingStrategy, model_path: str) -> LlamaFactoryConfig:
+    """Create a default configuration based on the training strategy."""
+    
+    if strategy == TrainingStrategy.SUPERVISED_FINETUNING:
+        return SFTConfig(model_path=model_path)
+    elif strategy == TrainingStrategy.REINFORCEMENT_LEARNING:
+        return RLConfig(model_path=model_path)
+    elif strategy == TrainingStrategy.PREFERENCE_OPTIMIZATION:
+        return DPOConfig(model_path=model_path)
+    else:
+        raise ValueError(f"Unsupported strategy: {strategy}") 

isa_model/training/engine/llama_factory/data_adapter.py

@@ -0,0 +1,284 @@
+"""
+Data adapters for LlamaFactory training.
+"""
+
+import os
+import json
+import logging
+from typing import Dict, List, Any, Optional, Union
+
+from .config import DatasetFormat
+
+logger = logging.getLogger(__name__)
+
+
+class DataAdapter:
+    """
+    Base class for adapting datasets to LlamaFactory format.
+    
+    This class handles converting data from various formats into
+    the specific format expected by LlamaFactory.
+    
+    Example:
+        ```python
+        # Create a custom adapter
+        adapter = CustomAdapter()
+        
+        # Convert data
+        converted_path = adapter.convert_data(
+            input_file="path/to/source_data.json",
+            output_file="path/to/converted_data.json"
+        )
+        ```
+    """
+    
+    def __init__(self, format_type: DatasetFormat = DatasetFormat.ALPACA):
+        """
+        Initialize the data adapter.
+        
+        Args:
+            format_type: The target format to convert to
+        """
+        self.format_type = format_type
+        
+    def convert_data(self, input_file: str, output_file: Optional[str] = None) -> str:
+        """
+        Convert data from input format to LlamaFactory format.
+        
+        Args:
+            input_file: Path to the input data file
+            output_file: Path to save the converted data
+            
+        Returns:
+            Path to the converted data file
+        """
+        if output_file is None:
+            dirname = os.path.dirname(input_file)
+            basename = os.path.basename(input_file)
+            name, ext = os.path.splitext(basename)
+            output_file = os.path.join(dirname, f"{name}_converted{ext}")
+            
+        data = self._load_data(input_file)
+        converted_data = self._convert_data(data)
+        self._save_data(converted_data, output_file)
+        
+        return output_file
+        
+    def _load_data(self, input_file: str) -> Any:
+        """
+        Load data from the input file.
+        
+        Args:
+            input_file: Path to the input file
+            
+        Returns:
+            Loaded data
+        """
+        with open(input_file, 'r', encoding='utf-8') as f:
+            return json.load(f)
+            
+    def _save_data(self, data: Any, output_file: str) -> None:
+        """
+        Save data to the output file.
+        
+        Args:
+            data: Data to save
+            output_file: Path to save the data
+        """
+        with open(output_file, 'w', encoding='utf-8') as f:
+            json.dump(data, f, ensure_ascii=False, indent=2)
+            
+    def _convert_data(self, data: Any) -> List[Dict[str, Any]]:
+        """
+        Convert data to the target format.
+        
+        Args:
+            data: Input data
+            
+        Returns:
+            Converted data
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+
+class AlpacaAdapter(DataAdapter):
+    """
+    Adapter for Alpaca format data.
+    
+    Example:
+        ```python
+        adapter = AlpacaAdapter()
+        converted_path = adapter.convert_data("custom_data.json")
+        ```
+    """
+    
+    def __init__(self):
+        """Initialize the Alpaca adapter."""
+        super().__init__(DatasetFormat.ALPACA)
+        
+    def _convert_data(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Convert data to Alpaca format.
+        
+        Expected output format:
+        [
+            {
+                "instruction": "Task description",
+                "input": "Optional input (context)",
+                "output": "Expected output"
+            },
+            ...
+        ]
+        
+        Args:
+            data: Input data
+            
+        Returns:
+            Converted data in Alpaca format
+        """
+        result = []
+        
+        for item in data:
+            if isinstance(item, dict):
+                # If already in the expected format, just add to result
+                if all(k in item for k in ["instruction", "output"]):
+                    alpaca_item = {
+                        "instruction": item["instruction"],
+                        "input": item.get("input", ""),
+                        "output": item["output"]
+                    }
+                    result.append(alpaca_item)
+                # Otherwise, try to convert from common formats
+                elif "prompt" in item and "response" in item:
+                    alpaca_item = {
+                        "instruction": item["prompt"],
+                        "input": "",
+                        "output": item["response"]
+                    }
+                    result.append(alpaca_item)
+                elif "question" in item and "answer" in item:
+                    alpaca_item = {
+                        "instruction": item["question"],
+                        "input": "",
+                        "output": item["answer"]
+                    }
+                    result.append(alpaca_item)
+                else:
+                    logger.warning(f"Could not convert item: {item}")
+        
+        return result
+
+
+class ShareGPTAdapter(DataAdapter):
+    """
+    Adapter for ShareGPT format data.
+    
+    Example:
+        ```python
+        adapter = ShareGPTAdapter()
+        converted_path = adapter.convert_data("sharegpt_data.json")
+        ```
+    """
+    
+    def __init__(self):
+        """Initialize the ShareGPT adapter."""
+        super().__init__(DatasetFormat.SHAREGPT)
+        
+    def _convert_data(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Convert data to ShareGPT format.
+        
+        Expected output format:
+        [
+            {
+                "conversations": [
+                    {"from": "human", "value": "Human message"},
+                    {"from": "gpt", "value": "Assistant response"},
+                    ...
+                ]
+            },
+            ...
+        ]
+        
+        Args:
+            data: Input data
+            
+        Returns:
+            Converted data in ShareGPT format
+        """
+        result = []
+        
+        for item in data:
+            if isinstance(item, dict):
+                conversations = []
+                
+                # Handle different input formats
+                
+                # If already in conversations format
+                if "conversations" in item and isinstance(item["conversations"], list):
+                    # Make sure format is correct
+                    for conv in item["conversations"]:
+                        if "from" in conv and "value" in conv:
+                            # Normalize role names
+                            role = conv["from"].lower()
+                            if role in ["user", "human"]:
+                                role = "human"
+                            elif role in ["assistant", "gpt", "ai"]:
+                                role = "gpt"
+                            else:
+                                logger.warning(f"Unknown role: {role}, skipping message")
+                                continue
+                                
+                            conversations.append({
+                                "from": role,
+                                "value": conv["value"]
+                            })
+                    
+                # If in QA format
+                elif "question" in item and "answer" in item:
+                    conversations = [
+                        {"from": "human", "value": item["question"]},
+                        {"from": "gpt", "value": item["answer"]}
+                    ]
+                
+                # If in prompt/response format
+                elif "prompt" in item and "response" in item:
+                    conversations = [
+                        {"from": "human", "value": item["prompt"]},
+                        {"from": "gpt", "value": item["response"]}
+                    ]
+                    
+                if conversations:
+                    result.append({"conversations": conversations})
+        
+        return result
+
+
+class DataAdapterFactory:
+    """
+    Factory for creating data adapters.
+    
+    Example:
+        ```python
+        # Create an adapter based on format
+        adapter = DataAdapterFactory.create_adapter(DatasetFormat.ALPACA)
+        ```
+    """
+    
+    @staticmethod
+    def create_adapter(format_type: DatasetFormat) -> DataAdapter:
+        """
+        Create a data adapter for the specified format.
+        
+        Args:
+            format_type: The dataset format
+            
+        Returns:
+            A data adapter instance
+        """
+        if format_type == DatasetFormat.ALPACA:
+            return AlpacaAdapter()
+        elif format_type == DatasetFormat.SHAREGPT:
+            return ShareGPTAdapter()
+        else:
+            raise ValueError(f"Unsupported format type: {format_type}") 

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

@@ -0,0 +1,6 @@
+"""
+Example scripts for LlamaFactory fine-tuning and reinforcement learning.
+
+This package contains example scripts that demonstrate how to use
+the LlamaFactory engine for various tasks such as fine-tuning and RLHF.
+"""