Upload inference.py

inference.py

@@ -0,0 +1,78 @@
+"""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
+
+# Same architecture as training
+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}")