| """Retail World Model Inference Script |
| Predicts future retail sales given historical context using the trained world model. |
| """ |
| import os, pickle, numpy as np, pandas as pd, torch, torch.nn as nn |
| from transformers import AutoModelForSeq2SeqLM, AutoConfig |
| from sklearn.preprocessing import StandardScaler, LabelEncoder |
|
|
| |
| 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): |
| 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) |
| return {'mean': mean, 'var': var} |
|
|
| def load_model_and_scaler(checkpoint_dir, base_model_name='amazon/chronos-bolt-small', |
| context_len=60, pred_len=14, num_variates=5, embed_dim=64): |
| model = RetailWorldModel(base_model_name, context_len, pred_len, num_variates, embed_dim) |
| ckpt = torch.load(os.path.join(checkpoint_dir, 'pytorch_model.bin'), map_location='cpu') |
| model.load_state_dict(ckpt, strict=False) |
| with open(os.path.join(checkpoint_dir, 'scaler.pkl'), 'rb') as f: |
| scaler = pickle.load(f) |
| return model, scaler |
|
|
| def predict(model, scaler, context_history, steps=14): |
| """ |
| context_history: numpy array (context_len, num_variates) - last 60 days |
| Returns: dict with 'mean' (actual sales), 'lower', 'upper' (90% CI) |
| """ |
| model.eval() |
| with torch.no_grad(): |
| ctx = torch.tensor(context_history).unsqueeze(0).float() |
| out = model(ctx) |
| mean = out['mean'].squeeze(0).numpy() |
| std = np.sqrt(out['var'].squeeze(0).numpy()) |
| mean_sales = scaler.inverse_transform(mean.reshape(-1, 1)).flatten() |
| std_sales = std * scaler.scale_[0] |
| return { |
| 'mean': mean_sales, |
| 'lower': mean_sales - 1.645 * std_sales, |
| 'upper': mean_sales + 1.645 * std_sales, |
| } |
|
|
| if __name__ == '__main__': |
| import argparse |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--checkpoint', required=True) |
| parser.add_argument('--history', required=True, help='CSV with 60 rows of history') |
| parser.add_argument('--output', default='predictions.csv') |
| args = parser.parse_args() |
|
|
| model, scaler = load_model_and_scaler(args.checkpoint) |
| hist = pd.read_csv(args.history) |
| pred = predict(model, scaler, hist.values) |
| df = pd.DataFrame({'day': range(1, len(pred['mean'])+1), |
| 'predicted_sales': pred['mean'], |
| 'lower_90': pred['lower'], |
| 'upper_90': pred['upper']}) |
| df.to_csv(args.output, index=False) |
| print(f"Saved predictions to {args.output}") |
|
|
