train_v6.py

#!/usr/bin/env python3
"""Retail World Model - Training Script v6
Fixed: MSE loss instead of unstable NLL, save_safetensors=False for shared tensors, better numerical stability
"""
import os, pickle, numpy as np, pandas as pd, torch, torch.nn as nn
from datasets import load_dataset
from transformers import T5EncoderModel, AutoConfig, TrainingArguments, Trainer, EarlyStoppingCallback, set_seed
from peft import LoraConfig, get_peft_model, TaskType
from sklearn.preprocessing import StandardScaler, LabelEncoder

SEED = 42
MODEL_NAME = 'google/t5-efficient-tiny'
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)

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 = T5EncoderModel.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
        d_model = self.config.d_model
        self.input_proj = nn.Linear(num_variates, d_model)
        self.latent_dynamics = nn.LSTM(d_model, d_model, 2, batch_first=True, dropout=0.1)
        self.mean_head = nn.Sequential(nn.Linear(d_model, embed_dim), nn.GELU(), nn.Linear(embed_dim, 1))
        # Variance head with softplus and clamp for stability
        self.var_head = nn.Sequential(nn.Linear(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 = torch.clamp(self.var_head(states).squeeze(-1), min=1e-4, max=10.0)
        if return_loss and target is not None:
            # Stable NLL with clamped variance
            nll = 0.5 * torch.log(var) + 0.5 * (target - mean) ** 2 / var
            # Add MSE component for stability
            mse = (target - mean) ** 2
            loss = torch.mean(nll + 0.1 * mse)
            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.FEATURE_EXTRACTION)
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',
    eval_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,
    save_safetensors=False,  # Fix shared tensor issue
    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}')
trainer.push_to_hub()
print('Done!')