File size: 6,249 Bytes
89ea727 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | """
TFT training and inference wrapper.
Uses Nixtla's neuralforecast library which follows the (unique_id, ds, y) convention.
The model is trained once and cached; the Streamlit app calls forecast() at runtime.
"""
import os
import pickle
from pathlib import Path
from typing import List, Optional
import numpy as np
import pandas as pd
from neuralforecast import NeuralForecast
from neuralforecast.models import TFT
from neuralforecast.losses.pytorch import MQLoss
from src.forecasting.config import TFTConfig, DEFAULT_CONFIG
# ── Feature engineering ───────────────────────────────────────────────────────
FASHION_WEEK_MONTHS = {(1, 15), (1, 20), (1, 25), (5, 10), (5, 18),
(9, 25), (10, 2), (10, 8), (10, 15)}
EVENT_MAP = {
"FashionWeek": 1, "ProductDrop": 2, "Christmas": 3,
"ValentinesDay": 4, "ChineseNewYear": 5, "MothersDaySurge": 6,
}
def prepare_features(df: pd.DataFrame) -> pd.DataFrame:
"""
Encode categorical and temporal features required by the TFT config.
Operates on the full long-format DataFrame.
"""
df = df.copy()
df["ds"] = pd.to_datetime(df["ds"])
# Temporal features
df["day_of_week"] = df["ds"].dt.dayofweek / 6.0 # [0, 1]
df["month"] = (df["ds"].dt.month - 1) / 11.0 # [0, 1]
df["event_flag"] = df["event_name"].map(EVENT_MAP).fillna(0).astype(float)
df["is_fashion_week"] = (
df[["ds"]].assign(
key=list(zip(df["ds"].dt.month, df["ds"].dt.day))
)["key"].isin(FASHION_WEEK_MONTHS)
).astype(float)
# Static categorical encodings (label-encode per series)
category_enc = {c: i for i, c in enumerate(df["category"].unique())}
store_enc = {s: i for i, s in enumerate(df["store"].unique())}
df["category_enc"] = df["category"].map(category_enc).astype(float)
df["store_enc"] = df["store"].map(store_enc).astype(float)
# Price normalisation (min-max per item)
df["price"] = df.groupby("unique_id")["price"].transform(
lambda x: (x - x.min()) / (x.max() - x.min() + 1e-8)
)
return df
def train(
df: pd.DataFrame,
cfg: TFTConfig = DEFAULT_CONFIG,
val_size: int = 56, # last 8 weeks as validation
) -> NeuralForecast:
"""
Train TFT on the full dataset and persist the model.
Returns the fitted NeuralForecast object.
"""
df = prepare_features(df)
model = TFT(
h=cfg.h,
input_size=cfg.input_size,
hidden_size=cfg.hidden_size,
n_head=cfg.n_head,
attn_dropout=cfg.attn_dropout,
dropout=cfg.dropout,
loss=MQLoss(level=[80, 95]),
learning_rate=cfg.learning_rate,
max_steps=cfg.max_steps,
batch_size=cfg.batch_size,
val_check_steps=cfg.val_check_steps,
stat_exog_list=cfg.stat_exog_list,
hist_exog_list=cfg.hist_exog_list,
futr_exog_list=cfg.futr_exog_list,
scaler_type="standard",
)
nf = NeuralForecast(models=[model], freq="D")
nf.fit(df=df, val_size=val_size)
Path(cfg.model_dir).mkdir(exist_ok=True)
nf.save(path=f"{cfg.model_dir}/{cfg.model_name}", overwrite=True)
return nf
def load_model(cfg: TFTConfig = DEFAULT_CONFIG) -> NeuralForecast:
"""Load a previously trained model from disk."""
return NeuralForecast.load(path=f"{cfg.model_dir}/{cfg.model_name}")
def forecast(
nf: NeuralForecast,
df: pd.DataFrame,
unique_ids: Optional[List[str]] = None,
cfg: TFTConfig = DEFAULT_CONFIG,
) -> pd.DataFrame:
"""
Run inference for the given series.
Returns a DataFrame with columns:
unique_id, ds, TFT (point), TFT-lo-80, TFT-hi-80, TFT-lo-95, TFT-hi-95
"""
df = prepare_features(df)
if unique_ids is not None:
df = df[df["unique_id"].isin(unique_ids)]
# Build future exogenous DataFrame for the forecast horizon
futr_df = _build_future_exog(df, cfg.h)
preds = nf.predict(df=df, futr_df=futr_df)
return preds.reset_index()
def _build_future_exog(df: pd.DataFrame, h: int) -> pd.DataFrame:
"""
Extend the known covariates (event_flag, day_of_week, month, is_fashion_week)
into the forecast horizon for each series.
"""
last_dates = df.groupby("unique_id")["ds"].max()
records = []
for uid, last_date in last_dates.items():
future_dates = pd.date_range(
start=last_date + pd.Timedelta(days=1), periods=h, freq="D"
)
fw_flags = pd.Series(future_dates).apply(
lambda d: float((d.month, d.day) in FASHION_WEEK_MONTHS)
)
records.append(pd.DataFrame({
"unique_id": uid,
"ds": future_dates,
"day_of_week": future_dates.dayofweek / 6.0,
"month": (future_dates.month - 1) / 11.0,
"event_flag": 0.0, # conservative: no known future events
"is_fashion_week": fw_flags.values,
}))
return pd.concat(records, ignore_index=True)
# ── Benchmark utilities ───────────────────────────────────────────────────────
def compute_metrics(actuals: pd.Series, predictions: pd.Series) -> dict:
"""MAE, RMSE, MASE (vs. seasonal naïve baseline with period=7)."""
mae = np.abs(actuals - predictions).mean()
rmse = np.sqrt(((actuals - predictions) ** 2).mean())
naive = np.abs(actuals.values[7:] - actuals.values[:-7]).mean() + 1e-8
mase = mae / naive
return {"MAE": round(mae, 2), "RMSE": round(rmse, 2), "MASE": round(mase, 3)}
# ── Entrypoint ────────────────────────────────────────────────────────────────
if __name__ == "__main__":
from src.data.loader import generate_dataset
print("Generating dataset...")
df = generate_dataset()
print(f" {len(df):,} rows | {df['unique_id'].nunique()} series")
print("Training TFT model...")
nf = train(df)
print(" Model saved to models/tft_luxury")
|