StockSignalAI / model_trainer.py
pareshmishra's picture
added new files
a9237fb verified
Raw History Blame Contribute Delete
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)