#!/usr/bin/env python3 import os, pickle, numpy as np, pandas as pd, torch, torch.nn as nn from datasets import load_dataset from transformers import AutoModelForSeq2SeqLM, AutoConfig, TrainingArguments, Trainer, EarlyStoppingCallback, set_seed from peft import LoraConfig, get_peft_model, TaskType import trackio from sklearn.preprocessing import StandardScaler, LabelEncoder SEED = 42 MODEL_NAME = 'amazon/chronos-bolt-small' OUTPUT_DIR = '/tmp/outputs' HUB_MODEL_ID = 'superdkj/retail-world-model-v1' DATASET_NAME = 't4tiana/store-sales-time-series-forecasting' CONTEXT_LENGTH = 60 PREDICTION_LENGTH = 14 NUM_VARIATES = 5 EMBED_DIM = 64 set_seed(SEED) trackio.init(project='retail-world-model', run_name='retail-world-model-v1') class RetailWorldModel(nn.Module): def __init__(self, base_model_name, context_len, pred_len, num_variates, embed_dim): super().__init__() self.config = AutoConfig.from_pretrained(base_model_name) self.encoder = AutoModelForSeq2SeqLM.from_pretrained(base_model_name) self.context_len = context_len self.pred_len = pred_len self.num_variates = num_variates self.embed_dim = embed_dim self.input_proj = nn.Linear(num_variates, self.config.d_model) self.latent_dynamics = nn.LSTM(self.config.d_model, self.config.d_model, 2, batch_first=True, dropout=0.1) self.mean_head = nn.Sequential(nn.Linear(self.config.d_model, embed_dim), nn.GELU(), nn.Linear(embed_dim, 1)) self.var_head = nn.Sequential(nn.Linear(self.config.d_model, embed_dim), nn.GELU(), nn.Linear(embed_dim, 1), nn.Softplus()) def forward(self, context, target=None, return_loss=True): x = self.input_proj(context) enc_out = self.encoder.encoder(inputs_embeds=x, return_dict=True).last_hidden_state h0 = enc_out[:, -1:, :].transpose(0, 1).repeat(2, 1, 1) c0 = torch.zeros_like(h0) states = [] curr = enc_out[:, -1:, :] for _ in range(self.pred_len): out, (h0, c0) = self.latent_dynamics(curr, (h0, c0)) states.append(out) curr = out states = torch.cat(states, dim=1) mean = self.mean_head(states).squeeze(-1) var = self.var_head(states).squeeze(-1) if return_loss and target is not None: loss = torch.mean(0.5 * torch.log(var + 1e-6) + 0.5 * (target - mean) ** 2 / (var + 1e-6)) return {'loss': loss, 'mean': mean, 'var': var} return {'mean': mean, 'var': var} class RetailDataset(torch.utils.data.Dataset): def __init__(self, df, context_len=60, pred_len=14, scaler=None, fit_scaler=False): self.context_len = context_len self.pred_len = pred_len df = df.copy() df['date'] = pd.to_datetime(df['date']) df['day_of_week'] = df['date'].dt.dayofweek / 6.0 df['month'] = df['date'].dt.month / 12.0 self.family_enc = LabelEncoder() df['family_enc'] = self.family_enc.fit_transform(df['family']) df['family_enc'] = df['family_enc'] / len(self.family_enc.classes_) self.groups = [] for _, g in df.groupby(['store_nbr', 'family']): g = g.sort_values('date').reset_index(drop=True) if len(g) >= context_len + pred_len: self.groups.append(g) if scaler is None: all_sales = np.concatenate([g['sales'].values for g in self.groups]) self.scaler = StandardScaler() self.scaler.fit(all_sales.reshape(-1, 1)) else: self.scaler = scaler for i, g in enumerate(self.groups): g = g.copy() g['sales_scaled'] = self.scaler.transform(g['sales'].values.reshape(-1, 1)).flatten() self.groups[i] = g self.windows = [] for g in self.groups: for start in range(0, len(g) - context_len - pred_len + 1, 7): end_ctx = start + context_len end_pred = end_ctx + pred_len ctx = g.iloc[start:end_ctx][['sales_scaled', 'onpromotion', 'day_of_week', 'month', 'family_enc']].values.astype(np.float32) tgt = g.iloc[end_ctx:end_pred]['sales_scaled'].values.astype(np.float32) self.windows.append((ctx, tgt)) def __len__(self): return len(self.windows) def __getitem__(self, idx): ctx, tgt = self.windows[idx] return {'context': torch.tensor(ctx), 'target': torch.tensor(tgt)} def collate_fn(batch): return {'context': torch.stack([b['context'] for b in batch]), 'target': torch.stack([b['target'] for b in batch])} class RetailTrainer(Trainer): def compute_loss(self, model, inputs, return_outputs=False, **kwargs): out = model(inputs['context'], inputs['target'], return_loss=True) loss = out['loss'] if return_outputs: return loss, out return loss def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None): with torch.no_grad(): out = model(inputs['context'], inputs['target'], return_loss=True) loss = out['loss'] if prediction_loss_only: return (loss, None, None) return (loss, out['mean'], inputs['target']) print('Loading dataset...') ds = load_dataset(DATASET_NAME, split='train') df = ds.to_pandas() print(f'Rows: {len(df)}, Stores: {df["store_nbr"].nunique()}, Families: {df["family"].nunique()}') df['date'] = pd.to_datetime(df['date']) split_date = df['date'].max() - pd.Timedelta(days=90) train_df = df[df['date'] <= split_date] val_df = df[df['date'] > split_date] print(f'Train: {len(train_df)}, Val: {len(val_df)}') print('Building datasets...') train_ds = RetailDataset(train_df, CONTEXT_LENGTH, PREDICTION_LENGTH, fit_scaler=True) val_ds = RetailDataset(val_df, CONTEXT_LENGTH, PREDICTION_LENGTH, scaler=train_ds.scaler, fit_scaler=False) print(f'Train windows: {len(train_ds)}, Val windows: {len(val_ds)}') os.makedirs(OUTPUT_DIR, exist_ok=True) scaler_path = os.path.join(OUTPUT_DIR, 'scaler.pkl') with open(scaler_path, 'wb') as f: pickle.dump(train_ds.scaler, f) print('Initializing model...') model = RetailWorldModel(MODEL_NAME, CONTEXT_LENGTH, PREDICTION_LENGTH, NUM_VARIATES, EMBED_DIM) lora_cfg = LoraConfig(r=16, lora_alpha=32, target_modules=['q', 'v', 'k', 'o'], lora_dropout=0.05, bias='none', task_type=TaskType.SEQ_2_SEQ_LM) model.encoder = get_peft_model(model.encoder, lora_cfg) model.encoder.print_trainable_parameters() args = TrainingArguments( output_dir=OUTPUT_DIR, num_train_epochs=10, per_device_train_batch_size=32, per_device_eval_batch_size=64, learning_rate=1e-4, weight_decay=0.01, warmup_ratio=0.1, lr_scheduler_type='cosine', evaluation_strategy='epoch', save_strategy='epoch', logging_strategy='steps', logging_steps=50, logging_first_step=True, disable_tqdm=True, load_best_model_at_end=True, metric_for_best_model='eval_loss', greater_is_better=False, push_to_hub=True, hub_model_id=HUB_MODEL_ID, hub_strategy='every_save', save_total_limit=2, report_to='trackio', run_name='retail-world-model-v1', seed=SEED, dataloader_num_workers=4, gradient_accumulation_steps=2, fp16=True, ) trainer = RetailTrainer( model=model, args=args, train_dataset=train_ds, eval_dataset=val_ds, data_collator=collate_fn, callbacks=[EarlyStoppingCallback(early_stopping_patience=3)], ) print('Training...') trainer.train() trainer.save_model(os.path.join(OUTPUT_DIR, 'final')) eval_results = trainer.evaluate() print(f'Final eval_loss: {eval_results["eval_loss"]:.4f}') trackio.alert(title='Training Complete', text=f'Final eval_loss={eval_results["eval_loss"]:.4f}', level='INFO') trainer.push_to_hub() print('Done!')