aiecs 1.1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

aiecs/tools/statistics/model_trainer_tool.py

@@ -0,0 +1,470 @@
+"""
+Model Trainer Tool - AutoML and machine learning model training
+
+This tool provides AutoML capabilities with:
+- Automatic model selection for classification and regression
+- Hyperparameter tuning
+- Model evaluation and comparison
+- Feature importance analysis
+- Model explanation support
+"""
+
+import logging
+from typing import Dict, Any, List, Optional, Union
+from enum import Enum
+import pickle
+
+import pandas as pd
+import numpy as np
+from sklearn.model_selection import train_test_split, cross_val_score
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, r2_score, mean_squared_error
+from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor, GradientBoostingClassifier, GradientBoostingRegressor
+from sklearn.linear_model import LogisticRegression, LinearRegression
+from sklearn.preprocessing import LabelEncoder
+from pydantic import BaseModel, Field, ValidationError, ConfigDict
+
+from aiecs.tools.base_tool import BaseTool
+from aiecs.tools import register_tool
+
+
+class ModelType(str, Enum):
+    """Supported model types"""
+    LOGISTIC_REGRESSION = "logistic_regression"
+    LINEAR_REGRESSION = "linear_regression"
+    RANDOM_FOREST_CLASSIFIER = "random_forest_classifier"
+    RANDOM_FOREST_REGRESSOR = "random_forest_regressor"
+    GRADIENT_BOOSTING_CLASSIFIER = "gradient_boosting_classifier"
+    GRADIENT_BOOSTING_REGRESSOR = "gradient_boosting_regressor"
+    AUTO = "auto"
+
+
+class TaskType(str, Enum):
+    """Machine learning task types"""
+    CLASSIFICATION = "classification"
+    REGRESSION = "regression"
+    CLUSTERING = "clustering"
+
+
+
+
+class ModelTrainerError(Exception):
+    """Base exception for ModelTrainer errors"""
+    pass
+
+
+class TrainingError(ModelTrainerError):
+    """Raised when model training fails"""
+    pass
+
+
+@register_tool('model_trainer')
+class ModelTrainerTool(BaseTool):
+    """
+    AutoML tool that can:
+    1. Train multiple model types
+    2. Perform hyperparameter tuning
+    3. Evaluate and compare models
+    4. Generate feature importance
+    5. Provide model explanations
+    """
+    
+    # Configuration schema
+    class Config(BaseModel):
+        """Configuration for the model trainer tool"""
+        model_config = ConfigDict(env_prefix="MODEL_TRAINER_")
+        
+        test_size: float = Field(
+            default=0.2,
+            description="Proportion of data to use for testing"
+        )
+        random_state: int = Field(
+            default=42,
+            description="Random state for reproducibility"
+        )
+        cv_folds: int = Field(
+            default=5,
+            description="Number of cross-validation folds"
+        )
+        enable_hyperparameter_tuning: bool = Field(
+            default=False,
+            description="Whether to enable hyperparameter tuning"
+        )
+        max_tuning_iterations: int = Field(
+            default=20,
+            description="Maximum number of hyperparameter tuning iterations"
+        )
+    
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        """Initialize ModelTrainerTool with settings"""
+        super().__init__(config)
+        
+        # Parse configuration
+        self.config = self.Config(**(config or {}))
+        
+        self.logger = logging.getLogger(__name__)
+        if not self.logger.handlers:
+            handler = logging.StreamHandler()
+            handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s'))
+            self.logger.addHandler(handler)
+        self.logger.setLevel(logging.INFO)
+        
+        self._init_external_tools()
+        self.trained_models = {}
+    
+    def _init_external_tools(self):
+        """Initialize external task tools"""
+        self.external_tools = {}
+    
+    # Schema definitions
+    class TrainModelSchema(BaseModel):
+        """Schema for train_model operation"""
+        data: Union[Dict[str, Any], List[Dict[str, Any]]] = Field(description="Training data")
+        target: str = Field(description="Target column name")
+        model_type: ModelType = Field(default=ModelType.AUTO, description="Model type to train")
+        auto_tune: bool = Field(default=False, description="Enable hyperparameter tuning")
+        cross_validation: int = Field(default=5, description="Number of CV folds")
+    
+    class AutoSelectModelSchema(BaseModel):
+        """Schema for auto_select_model operation"""
+        data: Union[Dict[str, Any], List[Dict[str, Any]]] = Field(description="Data for model selection")
+        target: str = Field(description="Target column name")
+        task_type: Optional[TaskType] = Field(default=None, description="Task type")
+    
+    class EvaluateModelSchema(BaseModel):
+        """Schema for evaluate_model operation"""
+        model_id: str = Field(description="ID of trained model")
+        test_data: Union[Dict[str, Any], List[Dict[str, Any]]] = Field(description="Test data")
+        target: str = Field(description="Target column name")
+    
+    def train_model(
+        self,
+        data: Union[Dict[str, Any], List[Dict[str, Any]], pd.DataFrame],
+        target: str,
+        model_type: ModelType = ModelType.AUTO,
+        auto_tune: bool = False,
+        cross_validation: int = 5
+    ) -> Dict[str, Any]:
+        """
+        Train and evaluate model.
+        
+        Args:
+            data: Training data
+            target: Target column name
+            model_type: Type of model to train (auto-selected if AUTO)
+            auto_tune: Enable hyperparameter tuning
+            cross_validation: Number of cross-validation folds
+            
+        Returns:
+            Dict containing:
+                - model_id: Unique identifier for trained model
+                - model_type: Type of model trained
+                - performance: Performance metrics
+                - feature_importance: Feature importance scores
+                - cross_validation_scores: CV scores
+        """
+        try:
+            df = self._to_dataframe(data)
+            
+            # Separate features and target
+            X = df.drop(columns=[target])
+            y = df[target]
+            
+            # Determine task type and model
+            task_type = self._determine_task_type(y)
+            
+            if model_type == ModelType.AUTO:
+                model_type = self._auto_select_model_type(task_type)
+                self.logger.info(f"Auto-selected model type: {model_type.value}")
+            
+            # Prepare data
+            X_processed, feature_names = self._preprocess_features(X)
+            
+            # Handle categorical target for classification
+            if task_type == TaskType.CLASSIFICATION:
+                label_encoder = LabelEncoder()
+                y = label_encoder.fit_transform(y)
+            else:
+                label_encoder = None
+            
+            # Split data
+            X_train, X_test, y_train, y_test = train_test_split(
+                X_processed, y,
+                test_size=self.config.test_size,
+                random_state=self.config.random_state
+            )
+            
+            # Create and train model
+            model = self._create_model(model_type)
+            model.fit(X_train, y_train)
+            
+            # Make predictions
+            y_pred = model.predict(X_test)
+            
+            # Calculate metrics
+            performance = self._calculate_metrics(y_test, y_pred, task_type)
+            
+            # Cross-validation
+            cv_scores = cross_val_score(model, X_processed, y, cv=cross_validation)
+            
+            # Feature importance
+            feature_importance = self._get_feature_importance(model, feature_names)
+            
+            # Store model
+            model_id = f"model_{len(self.trained_models) + 1}"
+            self.trained_models[model_id] = {
+                'model': model,
+                'model_type': model_type.value,
+                'task_type': task_type.value,
+                'feature_names': feature_names,
+                'label_encoder': label_encoder
+            }
+            
+            return {
+                'model_id': model_id,
+                'model_type': model_type.value,
+                'task_type': task_type.value,
+                'performance': performance,
+                'feature_importance': feature_importance,
+                'cross_validation_scores': {
+                    'scores': cv_scores.tolist(),
+                    'mean': float(cv_scores.mean()),
+                    'std': float(cv_scores.std())
+                },
+                'training_samples': len(X_train),
+                'test_samples': len(X_test)
+            }
+            
+        except Exception as e:
+            self.logger.error(f"Error training model: {e}")
+            raise TrainingError(f"Model training failed: {e}")
+    
+    def auto_select_model(
+        self,
+        data: Union[Dict[str, Any], List[Dict[str, Any]], pd.DataFrame],
+        target: str,
+        task_type: Optional[TaskType] = None
+    ) -> Dict[str, Any]:
+        """
+        Automatically select best model based on data characteristics.
+        
+        Args:
+            data: Data for model selection
+            target: Target column name
+            task_type: Optional task type (auto-determined if None)
+            
+        Returns:
+            Dict containing recommended model and reasoning
+        """
+        try:
+            df = self._to_dataframe(data)
+            y = df[target]
+            
+            # Determine task type
+            if task_type is None:
+                task_type = self._determine_task_type(y)
+            
+            # Select model
+            model_type = self._auto_select_model_type(task_type)
+            
+            # Provide reasoning
+            reasoning = self._explain_model_selection(df, y, task_type, model_type)
+            
+            return {
+                'recommended_model': model_type.value,
+                'task_type': task_type.value,
+                'reasoning': reasoning,
+                'confidence': 'high'
+            }
+            
+        except Exception as e:
+            self.logger.error(f"Error in auto model selection: {e}")
+            raise TrainingError(f"Model selection failed: {e}")
+    
+    def evaluate_model(
+        self,
+        model_id: str,
+        test_data: Union[Dict[str, Any], List[Dict[str, Any]], pd.DataFrame],
+        target: str
+    ) -> Dict[str, Any]:
+        """
+        Evaluate trained model on test data.
+        
+        Args:
+            model_id: ID of trained model
+            test_data: Test data
+            target: Target column name
+            
+        Returns:
+            Dict containing evaluation metrics
+        """
+        try:
+            if model_id not in self.trained_models:
+                raise TrainingError(f"Model {model_id} not found")
+            
+            df = self._to_dataframe(test_data)
+            X_test = df.drop(columns=[target])
+            y_test = df[target]
+            
+            model_info = self.trained_models[model_id]
+            model = model_info['model']
+            task_type = TaskType(model_info['task_type'])
+            
+            # Preprocess features
+            X_processed, _ = self._preprocess_features(X_test)
+            
+            # Handle label encoding for classification
+            if model_info['label_encoder']:
+                y_test = model_info['label_encoder'].transform(y_test)
+            
+            # Make predictions
+            y_pred = model.predict(X_processed)
+            
+            # Calculate metrics
+            performance = self._calculate_metrics(y_test, y_pred, task_type)
+            
+            return {
+                'model_id': model_id,
+                'performance': performance,
+                'test_samples': len(X_test)
+            }
+            
+        except Exception as e:
+            self.logger.error(f"Error evaluating model: {e}")
+            raise TrainingError(f"Model evaluation failed: {e}")
+    
+    def tune_hyperparameters(
+        self,
+        data: Union[Dict[str, Any], List[Dict[str, Any]], pd.DataFrame],
+        target: str,
+        model_type: ModelType
+    ) -> Dict[str, Any]:
+        """
+        Tune hyperparameters for specified model type.
+        
+        Args:
+            data: Training data
+            target: Target column name
+            model_type: Model type to tune
+            
+        Returns:
+            Dict containing best parameters and performance
+        """
+        try:
+            # Note: Full hyperparameter tuning with GridSearchCV would be implemented here
+            # For now, returning placeholder structure
+            self.logger.info("Hyperparameter tuning is a placeholder - train with default params")
+            
+            result = self.train_model(data, target, model_type, auto_tune=False)
+            result['tuning_note'] = "Using default parameters - full tuning not implemented"
+            
+            return result
+            
+        except Exception as e:
+            self.logger.error(f"Error tuning hyperparameters: {e}")
+            raise TrainingError(f"Hyperparameter tuning failed: {e}")
+    
+    # Internal helper methods
+    
+    def _to_dataframe(self, data: Union[Dict, List, pd.DataFrame]) -> pd.DataFrame:
+        """Convert data to DataFrame"""
+        if isinstance(data, pd.DataFrame):
+            return data
+        elif isinstance(data, list):
+            return pd.DataFrame(data)
+        elif isinstance(data, dict):
+            return pd.DataFrame([data])
+        else:
+            raise TrainingError(f"Unsupported data type: {type(data)}")
+    
+    def _determine_task_type(self, y: pd.Series) -> TaskType:
+        """Determine task type from target variable"""
+        if y.dtype in ['object', 'category', 'bool']:
+            return TaskType.CLASSIFICATION
+        elif y.nunique() < 10 and y.dtype in ['int64', 'int32']:
+            return TaskType.CLASSIFICATION
+        else:
+            return TaskType.REGRESSION
+    
+    def _auto_select_model_type(self, task_type: TaskType) -> ModelType:
+        """Auto-select model type based on task"""
+        if task_type == TaskType.CLASSIFICATION:
+            return ModelType.RANDOM_FOREST_CLASSIFIER
+        else:
+            return ModelType.RANDOM_FOREST_REGRESSOR
+    
+    def _create_model(self, model_type: ModelType):
+        """Create model instance"""
+        if model_type == ModelType.LOGISTIC_REGRESSION:
+            return LogisticRegression(random_state=self.config.random_state, max_iter=1000)
+        elif model_type == ModelType.LINEAR_REGRESSION:
+            return LinearRegression()
+        elif model_type == ModelType.RANDOM_FOREST_CLASSIFIER:
+            return RandomForestClassifier(random_state=self.config.random_state, n_estimators=100)
+        elif model_type == ModelType.RANDOM_FOREST_REGRESSOR:
+            return RandomForestRegressor(random_state=self.config.random_state, n_estimators=100)
+        elif model_type == ModelType.GRADIENT_BOOSTING_CLASSIFIER:
+            return GradientBoostingClassifier(random_state=self.config.random_state)
+        elif model_type == ModelType.GRADIENT_BOOSTING_REGRESSOR:
+            return GradientBoostingRegressor(random_state=self.config.random_state)
+        else:
+            raise TrainingError(f"Unsupported model type: {model_type}")
+    
+    def _preprocess_features(self, X: pd.DataFrame) -> tuple:
+        """Preprocess features for training"""
+        X_processed = X.copy()
+        
+        # Handle categorical variables with label encoding
+        for col in X_processed.select_dtypes(include=['object', 'category']).columns:
+            le = LabelEncoder()
+            X_processed[col] = le.fit_transform(X_processed[col].astype(str))
+        
+        # Handle missing values
+        X_processed = X_processed.fillna(X_processed.mean(numeric_only=True))
+        
+        feature_names = X_processed.columns.tolist()
+        
+        return X_processed.values, feature_names
+    
+    def _calculate_metrics(self, y_true, y_pred, task_type: TaskType) -> Dict[str, float]:
+        """Calculate performance metrics"""
+        if task_type == TaskType.CLASSIFICATION:
+            return {
+                'accuracy': float(accuracy_score(y_true, y_pred)),
+                'precision': float(precision_score(y_true, y_pred, average='weighted', zero_division=0)),
+                'recall': float(recall_score(y_true, y_pred, average='weighted', zero_division=0)),
+                'f1_score': float(f1_score(y_true, y_pred, average='weighted', zero_division=0))
+            }
+        else:
+            mse = mean_squared_error(y_true, y_pred)
+            return {
+                'r2_score': float(r2_score(y_true, y_pred)),
+                'mse': float(mse),
+                'rmse': float(np.sqrt(mse)),
+                'mae': float(np.mean(np.abs(y_true - y_pred)))
+            }
+    
+    def _get_feature_importance(self, model, feature_names: List[str]) -> Dict[str, float]:
+        """Extract feature importance from model"""
+        if hasattr(model, 'feature_importances_'):
+            importance = model.feature_importances_
+            return {name: float(imp) for name, imp in zip(feature_names, importance)}
+        elif hasattr(model, 'coef_'):
+            importance = np.abs(model.coef_).flatten()
+            return {name: float(imp) for name, imp in zip(feature_names, importance)}
+        else:
+            return {}
+    
+    def _explain_model_selection(self, df: pd.DataFrame, y: pd.Series, task_type: TaskType, model_type: ModelType) -> str:
+        """Explain why a model was selected"""
+        n_samples = len(df)
+        n_features = len(df.columns) - 1
+        
+        reasons = []
+        reasons.append(f"Task type: {task_type.value}")
+        reasons.append(f"Dataset size: {n_samples} samples, {n_features} features")
+        
+        if model_type in [ModelType.RANDOM_FOREST_CLASSIFIER, ModelType.RANDOM_FOREST_REGRESSOR]:
+            reasons.append("Random Forest selected for robust performance and feature importance")
+        
+        return "; ".join(reasons)
+