import os import warnings # Suppress deprecation and future warnings from third-party libraries (like MLflow and statsmodels) warnings.filterwarnings("ignore", category=FutureWarning) warnings.filterwarnings("ignore") import numpy as np import pandas as pd import mlflow import mlflow.xgboost import mlflow.statsmodels import xgboost as xgb from statsmodels.tsa.arima.model import ARIMA from sklearn.metrics import mean_squared_error, mean_absolute_error import logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') COMMODITY_MAP = { "Samba": 0, "Kekulu": 1, "Big Onion": 2, "Potato": 3, "Dried Chilli": 4, "Coconut": 5 } def load_real_data(): """Loads cleaned daily real price data.""" csv_path = "data/processed/clean_prices.csv" if not os.path.exists(csv_path): raise FileNotFoundError(f"Clean data not found at {csv_path}. Please run clean_data.py first.") df = pd.read_csv(csv_path) df['Date'] = pd.to_datetime(df['Date']) df.set_index('Date', inplace=True) return df def train_arima(train, test): """Trains an ARIMA baseline model and logs to MLflow.""" logging.info("Training ARIMA model...") with mlflow.start_run(run_name="ARIMA_Baseline"): # Define hyperparams dynamically to prevent crashing on short data ar_order = min(5, max(1, len(train) - 2)) order = (ar_order, 1, 0) mlflow.log_param("order", str(order)) # Train model = ARIMA(train, order=order) fitted_model = model.fit() # Predict predictions = fitted_model.forecast(steps=len(test)) # Metrics rmse = np.sqrt(mean_squared_error(test, predictions)) mae = mean_absolute_error(test, predictions) mlflow.log_metric("rmse", rmse) mlflow.log_metric("mae", mae) # Log model artifact mlflow.statsmodels.log_model(fitted_model, name="arima_model") logging.info(f"ARIMA RMSE: {rmse:.2f}, MAE: {mae:.2f}") def train_linear_regression(X_train, X_test, y_train, y_test): """Trains a baseline Linear Regression model on lag features for comparison and logs to MLflow.""" logging.info("Training Linear Regression baseline...") from sklearn.linear_model import LinearRegression with mlflow.start_run(run_name="LinearRegression_Baseline"): # Fit Linear Regression model model = LinearRegression() model.fit(X_train, y_train) # Predict on test set predictions = model.predict(X_test) # Calculate evaluation metrics rmse = np.sqrt(mean_squared_error(y_test, predictions)) mae = mean_absolute_error(y_test, predictions) # Log metrics to MLflow mlflow.log_metric("rmse", rmse) mlflow.log_metric("mae", mae) # Log model artifact mlflow.sklearn.log_model(model, name="linear_regression_model") logging.info(f"Linear Regression RMSE: {rmse:.2f}, MAE: {mae:.2f}") def train_xgboost(df): """Trains an XGBoost candidate model with lag features and logs to MLflow.""" logging.info("Training XGBoost model...") # Feature Engineering (Dynamic lag based on data size) min_days = df['Commodity'].value_counts().min() max_lags = max(1, min(7, min_days - 3)) if min_days <= 3: logging.warning("Data volume extremely low. Setting lag to 1. XGBoost model might not be robust.") max_lags = 1 df_list = [] for commodity in df['Commodity'].unique(): df_c = df[df['Commodity'] == commodity].copy() for i in range(1, max_lags + 1): df_c[f'Lag_{i}'] = df_c['Price'].shift(i) df_list.append(df_c) df_xgb = pd.concat(df_list) df_xgb.dropna(inplace=True) if len(df_xgb) < 2: logging.error("Not enough data to train XGBoost after lagging. Please gather more days of data.") return # Add Commodity_ID df_xgb['Commodity_ID'] = df_xgb['Commodity'].map(COMMODITY_MAP) feature_cols = ['Commodity_ID'] + [f'Lag_{i}' for i in range(1, max_lags + 1)] X = df_xgb[feature_cols] y = df_xgb['Price'] # Train-test split (Dynamic based on data size) test_size = max(1, int(len(X) * 0.2)) train_size = len(X) - test_size X_train, X_test = X.iloc[:train_size], X.iloc[train_size:] y_train, y_test = y.iloc[:train_size], y.iloc[train_size:] # Train Linear Regression baseline for academic performance comparison train_linear_regression(X_train, X_test, y_train, y_test) with mlflow.start_run(run_name="XGBoost_Candidate"): params = {"objective": "reg:squarederror", "n_estimators": 100, "learning_rate": 0.1, "max_depth": 5} mlflow.log_params(params) model = xgb.XGBRegressor(**params) model.fit(X_train, y_train) predictions = model.predict(X_test) rmse = np.sqrt(mean_squared_error(y_test, predictions)) mae = mean_absolute_error(y_test, predictions) mlflow.log_metric("rmse", rmse) mlflow.log_metric("mae", mae) # Log model artifact mlflow.xgboost.log_model(model, name="xgboost_model") logging.info(f"XGBoost RMSE: {rmse:.2f}, MAE: {mae:.2f}") if __name__ == "__main__": # Use environment variable for MLflow tracking URI, or default to local folder tracking_uri = os.environ.get("MLFLOW_TRACKING_URI", "file:./mlruns") mlflow.set_tracking_uri(tracking_uri) # If using local tracking, ensure the mlruns directory exists if tracking_uri.startswith("file:"): os.makedirs(tracking_uri.replace("file:", ""), exist_ok=True) # Configure DagsHub credentials automatically if URI is provided dagshub_token = os.environ.get("DAGSHUB_USER_TOKEN") if dagshub_token and "dagshub.com" in tracking_uri: # Extract username from URI (e.g. https://dagshub.com/username/repo.mlflow) username = tracking_uri.split('/')[-2] os.environ["MLFLOW_TRACKING_USERNAME"] = username os.environ["MLFLOW_TRACKING_PASSWORD"] = dagshub_token logging.info(f"Configured DagsHub MLflow remote for user: {username}") mlflow.set_experiment("SL_Commodity_Forecasting") df = load_real_data() # Split for ARIMA (Baseline uses only Samba) df_samba = df[df['Commodity'] == 'Samba'] if len(df_samba) < 3: logging.error("Not enough data for Samba to train ARIMA. Exiting.") else: test_size = max(1, int(len(df_samba) * 0.2)) train_size = len(df_samba) - test_size train, test = df_samba.iloc[:train_size]['Price'], df_samba.iloc[train_size:]['Price'] train_arima(train, test) train_xgboost(df) logging.info("Training complete. Run 'mlflow ui' in the terminal to view your experiment dashboards!")