Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| import torch.nn as nn | |
| from torch.utils.data import DataLoader, TensorDataset | |
| from prophet import Prophet | |
| from pmdarima import auto_arima | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| class LSTMForecaster(nn.Module): | |
| """LSTM model for time series forecasting.""" | |
| def __init__(self, input_size: int, hidden_size: int, num_layers: int, | |
| output_size: int, dropout: float = 0.2): | |
| super(LSTMForecaster, self).__init__() | |
| self.hidden_size = hidden_size | |
| self.num_layers = num_layers | |
| self.lstm = nn.LSTM(input_size, hidden_size, num_layers, | |
| batch_first=True, dropout=dropout) | |
| self.dropout = nn.Dropout(dropout) | |
| self.linear = nn.Linear(hidden_size, output_size) | |
| def forward(self, x): | |
| lstm_out, _ = self.lstm(x) | |
| lstm_out = self.dropout(lstm_out[:, -1, :]) # Take the last output | |
| out = self.linear(lstm_out) | |
| return out | |
| class AdvancedModelTrainer: | |
| """Trainer for advanced forecasting models.""" | |
| def __init__(self, config: dict): | |
| self.config = config | |
| def train_lstm(self, X_train: np.ndarray, y_train: np.ndarray, | |
| X_val: np.ndarray = None, | |
| y_val: np.ndarray = None) -> nn.Module: | |
| """Train LSTM model.""" | |
| model_config = self.config['lstm'] | |
| # Convert to PyTorch tensors | |
| train_dataset = TensorDataset( | |
| torch.FloatTensor(X_train), | |
| torch.FloatTensor(y_train) | |
| ) | |
| train_loader = DataLoader(train_dataset, batch_size=model_config['batch_size'], shuffle=True) | |
| # Initialize model | |
| model = LSTMForecaster( | |
| input_size=X_train.shape[2], | |
| hidden_size=model_config['hidden_size'], | |
| num_layers=model_config['num_layers'], | |
| output_size=y_train.shape[1], | |
| dropout=model_config['dropout'] | |
| ) | |
| # Training setup | |
| criterion = nn.MSELoss() | |
| optimizer = torch.optim.Adam(model.parameters(), lr=model_config['learning_rate']) | |
| # Training loop | |
| for epoch in range(model_config['epochs']): | |
| model.train() | |
| epoch_loss = 0 | |
| for batch_X, batch_y in train_loader: | |
| optimizer.zero_grad() | |
| predictions = model(batch_X) | |
| loss = criterion(predictions, batch_y) | |
| loss.backward() | |
| # Gradient clipping | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) | |
| optimizer.step() | |
| epoch_loss += loss.item() | |
| if epoch % 10 == 0: | |
| logger.info(f'Epoch {epoch}, Loss: {epoch_loss/len(train_loader):.4f}') | |
| return model | |
| def train_prophet(self, df: pd.DataFrame, | |
| date_col: str, | |
| value_col: str) -> Prophet: | |
| """Train Facebook Prophet model.""" | |
| prophet_df = df[[date_col, value_col]].rename( | |
| columns={date_col: 'ds', value_col: 'y'} | |
| ) | |
| model = Prophet( | |
| changepoint_prior_scale=self.config['prophet'].get('changepoint_prior_scale', 0.05), | |
| seasonality_prior_scale=self.config['prophet'].get('seasonality_prior_scale', 10), | |
| yearly_seasonality=self.config['prophet'].get('yearly_seasonality', True), | |
| weekly_seasonality=self.config['prophet'].get('weekly_seasonality', True), | |
| daily_seasonality=self.config['prophet'].get('daily_seasonality', False) | |
| ) | |
| model.fit(prophet_df) | |
| return model | |
| def train_auto_arima(self, series: pd.Series) -> object: | |
| """Train auto ARIMA model.""" | |
| model = auto_arima( | |
| series, | |
| start_p=1, | |
| start_q=1, | |
| max_p=3, | |
| max_q=3, | |
| seasonal=True, | |
| m=7, | |
| stepwise=True, | |
| suppress_warnings=True, | |
| error_action='ignore' | |
| ) | |
| return model |