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

isa_model/scripts/training_tracker.py

@@ -1,257 +0,0 @@
-"""
-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/__init__.py

@@ -1,74 +0,0 @@
-"""
-ISA Model Training Module
-
-Provides unified training capabilities for AI models including:
-- Local training with SFT (Supervised Fine-Tuning)
-- Cloud training on RunPod
-- Model evaluation and management
-- HuggingFace integration
-
-Example usage:
-    ```python
-    from isa_model.training import TrainingFactory, train_gemma
-    
-    # Quick Gemma training
-    model_path = train_gemma(
-        dataset_path="tatsu-lab/alpaca",
-        model_size="4b",
-        num_epochs=3
-    )
-    
-    # Advanced training with custom configuration
-    factory = TrainingFactory()
-    model_path = factory.train_model(
-        model_name="google/gemma-2-4b-it",
-        dataset_path="your-dataset.json",
-        use_lora=True,
-        batch_size=4,
-        num_epochs=3
-    )
-    ```
-"""
-
-# Import the new clean factory
-from .factory import TrainingFactory, train_gemma
-
-# Import core components
-from .core import (
-    TrainingConfig,
-    LoRAConfig, 
-    DatasetConfig,
-    BaseTrainer,
-    SFTTrainer,
-    TrainingUtils,
-    DatasetManager
-)
-
-# Import cloud training components
-from .cloud import (
-    RunPodConfig,
-    StorageConfig,
-    JobConfig,
-    TrainingJobOrchestrator
-)
-
-__all__ = [
-    # Main factory
-    'TrainingFactory',
-    'train_gemma',
-    
-    # Core components
-    'TrainingConfig',
-    'LoRAConfig',
-    'DatasetConfig', 
-    'BaseTrainer',
-    'SFTTrainer',
-    'TrainingUtils',
-    'DatasetManager',
-    
-    # Cloud components
-    'RunPodConfig',
-    'StorageConfig',
-    'JobConfig',
-    'TrainingJobOrchestrator'
-] 

isa_model/training/annotation/annotation_schema.py

@@ -1,47 +0,0 @@
-# app/services/llm_model/tracing/annotation/annotation_schema.py
-from enum import Enum
-from pydantic import BaseModel, Field
-from typing import Dict, Any, List, Optional
-from datetime import datetime
-
-class AnnotationType(str, Enum):
-    ACCURACY = "accuracy"
-    HELPFULNESS = "helpfulness"
-    TOXICITY = "toxicity"
-    CUSTOM = "custom"
-
-class RatingScale(int, Enum):
-    POOR = 1
-    FAIR = 2
-    GOOD = 3
-    EXCELLENT = 4
-
-class AnnotationAspects(BaseModel):
-    factually_correct: bool = True
-    relevant: bool = True
-    harmful: bool = False
-    biased: bool = False
-    complete: bool = True
-    efficient: bool = True
-
-class BetterResponse(BaseModel):
-    content: str
-    reason: Optional[str]
-    metadata: Optional[Dict[str, Any]] = {}
-
-class AnnotationFeedback(BaseModel):
-    rating: RatingScale
-    category: AnnotationType
-    aspects: AnnotationAspects
-    better_response: Optional[BetterResponse]
-    comment: Optional[str]
-    metadata: Optional[Dict[str, Any]] = {}
-    is_selected_for_training: bool = False  
-
-class ItemAnnotation(BaseModel):
-    item_id: str
-    feedback: Optional[AnnotationFeedback]
-    status: str = "pending"  
-    annotated_at: Optional[datetime]
-    annotator_id: Optional[str]
-    training_status: Optional[str] = None  

isa_model/training/annotation/processors/annotation_processor.py

@@ -1,126 +0,0 @@
-from typing import Dict, Any, List
-from datetime import datetime
-from app.config.config_manager import config_manager
-from app.services.training.llm_model.annotation.annotation_schema import AnnotationFeedback, RatingScale, AnnotationAspects
-from bson.objectid import ObjectId
-from app.services.training.llm_model.annotation.storage.dataset_manager import DatasetManager
-
-class AnnotationProcessor:
-    def __init__(self):
-        self.logger = config_manager.get_logger(__name__)
-        self.dataset_manager = DatasetManager()
-        self.batch_size = 1000  # Configure as needed
-    
-    async def process_queue(self) -> None:
-        """Process pending items and create datasets"""
-        db = await config_manager.get_db('mongodb')
-        queue = db['training_queue']
-        
-        # Process SFT items
-        sft_items = await self._get_pending_items("sft")
-        if len(sft_items) >= self.batch_size:
-            await self._create_sft_dataset(sft_items)
-        
-        # Process RLHF items
-        rlhf_items = await self._get_pending_items("rlhf")
-        if len(rlhf_items) >= self.batch_size:
-            await self._create_rlhf_dataset(rlhf_items)
-
-    async def _create_sft_dataset(self, items: List[Dict[str, Any]]):
-        """Create and upload SFT dataset"""
-        dataset = await self.dataset_manager.create_dataset(
-            name=f"sft_dataset_v{datetime.now().strftime('%Y%m%d')}",
-            type="sft",
-            version=datetime.now().strftime("%Y%m%d"),
-            source_annotations=[item["annotation_id"] for item in items]
-        )
-        
-        formatted_data = [
-            await self._process_sft_item(item) 
-            for item in items
-        ]
-        
-        await self.dataset_manager.upload_dataset_file(
-            dataset.id,
-            formatted_data
-        )
-
-    async def _process_sft_item(self, item: Dict[str, Any]) -> Dict[str, Any]:
-        """Process item for SFT dataset generation
-        Format follows HF conversation format for SFT training
-        """
-        db = await config_manager.get_db('mongodb')
-        annotations = db['annotations']
-        
-        # Get full annotation context
-        annotation = await annotations.find_one({"_id": ObjectId(item["annotation_id"])})
-        target_item = next(i for i in annotation["items"] if i["item_id"] == item["item_id"])
-        
-        # Format as conversation
-        messages = [
-            {
-                "role": "system",
-                "content": "You are a helpful AI assistant that provides accurate and relevant information."
-            },
-            {
-                "role": "user",
-                "content": target_item["input"]["messages"][0]["content"]
-            },
-            {
-                "role": "assistant", 
-                "content": target_item["output"]["content"]
-            }
-        ]
-
-        return {
-            "messages": messages,
-            "metadata": {
-                "rating": item["feedback"]["rating"],
-                "aspects": item["feedback"]["aspects"],
-                "category": item["feedback"]["category"]
-            }
-        }
-
-    async def _process_rlhf_item(self, item: Dict[str, Any]) -> Dict[str, Any]:
-        """Process item for RLHF dataset generation
-        Format follows preference pairs structure for RLHF training
-        """
-        db = await config_manager.get_db('mongodb')
-        annotations = db['annotations']
-        
-        # Get full annotation context
-        annotation = await annotations.find_one({"_id": ObjectId(item["annotation_id"])})
-        target_item = next(i for i in annotation["items"] if i["item_id"] == item["item_id"])
-        
-        # Format as preference pairs
-        return {
-            "prompt": target_item["input"]["messages"][0]["content"],
-            "chosen": item["feedback"]["better_response"]["content"],
-            "rejected": target_item["output"]["content"],
-            "metadata": {
-                "reason": item["feedback"]["better_response"]["reason"],
-                "category": item["feedback"]["category"]
-            }
-        }
-
-    async def get_training_data(
-        self,
-        data_type: str,
-        limit: int = 1000
-    ) -> List[Dict[str, Any]]:
-        """Retrieve formatted training data"""
-        db = await config_manager.get_db('mongodb')
-        training_data = db['training_data']
-        
-        data = await training_data.find(
-            {"type": data_type}
-        ).limit(limit).to_list(length=limit)
-        
-        if data_type == "sft":
-            return [item["data"]["messages"] for item in data]
-        else:  # rlhf
-            return [{
-                "prompt": item["data"]["prompt"],
-                "chosen": item["data"]["chosen"],
-                "rejected": item["data"]["rejected"]
-            } for item in data]

isa_model/training/annotation/storage/dataset_manager.py

@@ -1,131 +0,0 @@
-# app/services/llm_model/annotation/dataset/dataset_manager.py
-from typing import Dict, Any, List
-from datetime import datetime
-import json
-import io
-from app.config.config_manager import config_manager
-from .dataset_schema import Dataset, DatasetType, DatasetStatus, DatasetFiles, DatasetStats
-from bson import ObjectId
-
-class DatasetManager:
-    def __init__(self):
-        self.logger = config_manager.get_logger(__name__)
-        self.minio_client = None
-        self.bucket_name = "training-datasets"
-
-    async def _ensure_minio_client(self):
-        if not self.minio_client:
-            self.minio_client = await config_manager.get_storage_client()
-
-    async def create_dataset(
-        self, 
-        name: str,
-        type: DatasetType,
-        version: str,
-        source_annotations: List[str]
-    ) -> Dataset:
-        """Create a new dataset record"""
-        db = await config_manager.get_db('mongodb')
-        collection = db['training_datasets']
-        
-        dataset = Dataset(
-            name=name,
-            type=type,
-            version=version,
-            storage_path=f"datasets/{type.value}/{version}",
-            files=DatasetFiles(
-                train="train.jsonl",
-                eval=None,
-                test=None
-            ),
-            stats=DatasetStats(
-                total_examples=0,
-                avg_length=0.0,
-                num_conversations=0,
-                additional_metrics={}
-            ),
-            source_annotations=source_annotations,
-            created_at=datetime.utcnow(),
-            status=DatasetStatus.PENDING,
-            metadata={}
-        )
-        
-        result = await collection.insert_one(dataset.dict(exclude={'id'}))
-        return Dataset(**{**dataset.dict(), '_id': result.inserted_id})
-
-    async def upload_dataset_file(
-        self,
-        dataset_id: str,
-        data: List[Dict[str, Any]],
-        file_type: str = "train"
-    ) -> bool:
-        """Upload dataset to MinIO"""
-        try:
-            await self._ensure_minio_client()
-            db = await config_manager.get_db('mongodb')
-            
-            object_id = ObjectId(dataset_id)
-            dataset = await db['training_datasets'].find_one({"_id": object_id})
-            
-            if not dataset:
-                self.logger.error(f"Dataset not found with id: {dataset_id}")
-                return False
-            
-            # Convert to JSONL
-            buffer = io.StringIO()
-            for item in data:
-                buffer.write(json.dumps(item) + "\n")
-            
-            storage_path = dataset['storage_path'].rstrip('/')
-            file_path = f"{storage_path}/{file_type}.jsonl"
-            
-            buffer_value = buffer.getvalue().encode()
-            
-            self.logger.debug(f"Uploading to MinIO path: {file_path}")
-            
-            self.minio_client.put_object(
-                self.bucket_name,
-                file_path,
-                io.BytesIO(buffer_value),
-                len(buffer_value)
-            )
-            
-            avg_length = sum(len(str(item)) for item in data) / len(data) if data else 0
-            
-            await db['training_datasets'].update_one(
-                {"_id": object_id},
-                {
-                    "$set": {
-                        f"files.{file_type}": f"{file_type}.jsonl",
-                        "stats.total_examples": len(data),
-                        "stats.avg_length": avg_length,
-                        "stats.num_conversations": len(data),
-                        "status": DatasetStatus.READY
-                    }
-                }
-            )
-            
-            return True
-            
-        except Exception as e:
-            self.logger.error(f"Failed to upload dataset: {e}")
-            return False
-
-    async def get_dataset_info(self, dataset_id: str) -> Dict[str, Any]:
-        """Get dataset information"""
-        try:
-            db = await config_manager.get_db('mongodb')
-            object_id = ObjectId(dataset_id)  # Convert string ID to ObjectId
-            dataset = await db['training_datasets'].find_one({"_id": object_id})
-            
-            if not dataset:
-                self.logger.error(f"Dataset not found with id: {dataset_id}")
-                return None
-            
-            # Convert ObjectId to string for JSON serialization
-            dataset['_id'] = str(dataset['_id'])
-            return dataset
-            
-        except Exception as e:
-            self.logger.error(f"Failed to get dataset info: {e}")
-            return None

isa_model/training/annotation/storage/dataset_schema.py

@@ -1,44 +0,0 @@
-# app/services/llm_model/annotation/dataset/dataset_schema.py
-from enum import Enum
-from pydantic import BaseModel, Field
-from typing import Dict, List, Optional
-from datetime import datetime
-from bson import ObjectId
-
-class DatasetType(str, Enum):
-    SFT = "sft"
-    RLHF = "rlhf"
-
-class DatasetStatus(str, Enum):
-    PENDING = "pending"
-    PROCESSING = "processing"
-    READY = "ready"
-    ERROR = "error"
-
-class DatasetFiles(BaseModel):
-    train: str
-    eval: Optional[str]
-    test: Optional[str]
-
-class DatasetStats(BaseModel):
-    total_examples: int
-    avg_length: Optional[float]
-    num_conversations: Optional[int]
-    additional_metrics: Optional[Dict] = {}
-
-class Dataset(BaseModel):
-    id: Optional[ObjectId] = Field(None, alias="_id")
-    name: str
-    type: DatasetType
-    version: str
-    storage_path: str
-    files: DatasetFiles
-    stats: DatasetStats
-    source_annotations: List[str]
-    created_at: datetime
-    status: DatasetStatus
-    metadata: Optional[Dict] = {}
-
-    class Config:
-        arbitrary_types_allowed = True
-        populate_by_name = True