Spaces:
Sleeping
Sleeping
Download model_trainer.py from pareshmishra/StockSignalAI: direct link, hf CLI and curl.
- Browser
- Download file 7.19 kB
-
https://huggingface.co/spaces/pareshmishra/StockSignalAI/resolve/main/model_trainer.py
- Command line
-
hf download hf://spaces/pareshmishra/StockSignalAI/model_trainer.py
-
curl -L -o model_trainer.py https://huggingface.co/spaces/pareshmishra/StockSignalAI/resolve/main/model_trainer.py
7.19 kB
| import pandas as pd | |
| import numpy as np | |
| from sklearn.ensemble import RandomForestClassifier | |
| from sklearn.model_selection import train_test_split, cross_val_score, GridSearchCV | |
| from sklearn.metrics import classification_report, confusion_matrix, accuracy_score, precision_score, recall_score | |
| from sklearn.preprocessing import StandardScaler | |
| import xgboost as xgb | |
| import warnings | |
| warnings.filterwarnings('ignore') | |
| class ModelTrainer: | |
| def __init__(self): | |
| self.scaler = StandardScaler() | |
| def train_model(self, features, labels, model_type='random_forest'): | |
| """Train machine learning model for trading signals""" | |
| try: | |
| if len(features) == 0 or len(labels) == 0: | |
| raise ValueError("No data available for training") | |
| # Handle infinite and NaN values | |
| features = features.replace([np.inf, -np.inf], np.nan) | |
| features = features.fillna(features.mean()) | |
| # Split data | |
| X_train, X_test, y_train, y_test = train_test_split( | |
| features, labels, test_size=0.2, random_state=42, stratify=labels | |
| ) | |
| # Scale features | |
| X_train_scaled = self.scaler.fit_transform(X_train) | |
| X_test_scaled = self.scaler.transform(X_test) | |
| # Train model | |
| if model_type == 'random_forest': | |
| model = self.train_random_forest(X_train_scaled, y_train) | |
| elif model_type == 'xgboost': | |
| model = self.train_xgboost(X_train_scaled, y_train) | |
| else: | |
| raise ValueError(f"Unsupported model type: {model_type}") | |
| # Evaluate model | |
| metrics = self.evaluate_model(model, X_test_scaled, y_test) | |
| # Feature importance | |
| feature_importance = self.get_feature_importance(model, features.columns) | |
| # Create final model pipeline that includes scaling | |
| final_model = ModelPipeline(self.scaler, model) | |
| return final_model, features.columns.tolist(), metrics | |
| except Exception as e: | |
| print(f"Error training model: {str(e)}") | |
| raise e | |
| def train_random_forest(self, X_train, y_train): | |
| """Train Random Forest model with hyperparameter tuning""" | |
| try: | |
| # Define parameter grid for tuning | |
| param_grid = { | |
| 'n_estimators': [100, 200], | |
| 'max_depth': [10, 15, None], | |
| 'min_samples_split': [2, 5], | |
| 'min_samples_leaf': [1, 2] | |
| } | |
| # Create base model | |
| rf = RandomForestClassifier(random_state=42, class_weight='balanced') | |
| # Grid search with cross-validation | |
| grid_search = GridSearchCV( | |
| rf, param_grid, cv=5, scoring='accuracy', n_jobs=-1 | |
| ) | |
| grid_search.fit(X_train, y_train) | |
| return grid_search.best_estimator_ | |
| except Exception as e: | |
| print(f"Error training Random Forest: {str(e)}") | |
| # Fallback to basic model | |
| return RandomForestClassifier( | |
| n_estimators=100, random_state=42, class_weight='balanced' | |
| ).fit(X_train, y_train) | |
| def train_xgboost(self, X_train, y_train): | |
| """Train XGBoost model""" | |
| try: | |
| # Convert labels to ensure compatibility | |
| label_mapping = {-1: 0, 0: 1, 1: 2} | |
| y_train_mapped = np.array([label_mapping[label] for label in y_train]) | |
| model = xgb.XGBClassifier( | |
| n_estimators=100, | |
| max_depth=6, | |
| learning_rate=0.1, | |
| random_state=42, | |
| eval_metric='mlogloss' | |
| ) | |
| model.fit(X_train, y_train_mapped) | |
| # Wrap model to handle label mapping | |
| return XGBoostWrapper(model, label_mapping) | |
| except Exception as e: | |
| print(f"Error training XGBoost: {str(e)}") | |
| raise e | |
| def evaluate_model(self, model, X_test, y_test): | |
| """Evaluate model performance""" | |
| try: | |
| y_pred = model.predict(X_test) | |
| # Calculate metrics | |
| accuracy = accuracy_score(y_test, y_pred) | |
| precision = precision_score(y_test, y_pred, average='weighted', zero_division=0) | |
| recall = recall_score(y_test, y_pred, average='weighted', zero_division=0) | |
| metrics = { | |
| 'accuracy': accuracy, | |
| 'precision': precision, | |
| 'recall': recall, | |
| 'classification_report': classification_report(y_test, y_pred, zero_division=0), | |
| 'confusion_matrix': confusion_matrix(y_test, y_pred).tolist() | |
| } | |
| return metrics | |
| except Exception as e: | |
| print(f"Error evaluating model: {str(e)}") | |
| return {'accuracy': 0.0, 'precision': 0.0, 'recall': 0.0} | |
| def get_feature_importance(self, model, feature_names): | |
| """Get feature importance from trained model""" | |
| try: | |
| if hasattr(model, 'feature_importances_'): | |
| importance = model.feature_importances_ | |
| elif hasattr(model, 'model') and hasattr(model.model, 'feature_importances_'): | |
| importance = model.model.feature_importances_ | |
| else: | |
| return {} | |
| feature_importance = dict(zip(feature_names, importance)) | |
| # Sort by importance | |
| sorted_importance = dict(sorted(feature_importance.items(), | |
| key=lambda x: x[1], reverse=True)) | |
| return sorted_importance | |
| except Exception as e: | |
| print(f"Error getting feature importance: {str(e)}") | |
| return {} | |
| class ModelPipeline: | |
| """Pipeline that combines scaling and model prediction""" | |
| def __init__(self, scaler, model): | |
| self.scaler = scaler | |
| self.model = model | |
| def predict(self, X): | |
| """Make predictions with scaling""" | |
| X_scaled = self.scaler.transform(X) | |
| return self.model.predict(X_scaled) | |
| def predict_proba(self, X): | |
| """Make probability predictions with scaling""" | |
| X_scaled = self.scaler.transform(X) | |
| return self.model.predict_proba(X_scaled) | |
| class XGBoostWrapper: | |
| """Wrapper for XGBoost to handle label mapping""" | |
| def __init__(self, model, label_mapping): | |
| self.model = model | |
| self.label_mapping = label_mapping | |
| self.reverse_mapping = {v: k for k, v in label_mapping.items()} | |
| def predict(self, X): | |
| """Predict with label remapping""" | |
| pred_mapped = self.model.predict(X) | |
| return np.array([self.reverse_mapping[pred] for pred in pred_mapped]) | |
| def predict_proba(self, X): | |
| """Predict probabilities""" | |
| return self.model.predict_proba(X) | |