experiment-2107/scripts/model_transformers.py

"""Train a figure-skating action classifier on the pipeline's (T, 94)-feature tensors.

Variant of model.py: adds sinusoidal positional encoding before a real Transformer encoder
(self-attn + FFN per layer, via nn.TransformerEncoderLayer -- not a bare attention block).
Architecture: a Conv1D residual backbone, positional encoding, a 3-layer Transformer
encoder, temporal pooling, and a deep GELU dense head. Without positional encoding the
encoder is permutation-invariant over time (order only leaks in through the conv backbone's
local receptive field); this adds an explicit position signal.
Kept deliberately simple on the training-loop side (Adam + CrossEntropy + val-based best-model
selection). Metrics (precision / recall / F1) come from scikit-learn.

LABEL CONSISTENCY (the thing that must not break):
    The pipeline saves integer labels that are *exactly* labels.TAXONOMY indices (0..27).
    This model's output layer has num_classes = len(TAXONOMY) units, so output unit i is, by
    construction, the score for TAXONOMY[i]. No compaction / remapping is done anywhere, so
    argmax(logits) is directly a TAXONOMY index. Assertions below enforce this before training.

Usage:
    python model_transformers.py                 # train on /data/processed
    python model_transformers.py --data-dir DIR   # train on another processed dir
    python model_transformers.py --smoke          # self-test the plumbing on synthetic data
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import f1_score
from sklearn.utils.class_weight import compute_class_weight
from torch.utils.data import DataLoader, TensorDataset

try:
    from . import labels as labels_mod
except ImportError:
    import labels as labels_mod

from model import (
    ConvBlock,
    DenseBlock,
    load_split,
    resolve_taxonomy,
    assert_label_consistency,
    report_metrics,
    run_epoch,
    predict,
)


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

DEFAULT_DATA_DIR = "/data/processed"
EPOCHS = 100
BATCH_SIZE = 64
LEARNING_RATE = 5e-4       # lowered from 1e-3: the raw velocity features are heavy-tailed (|v|~960)
WEIGHT_DECAY = 1e-4
GRAD_CLIP = 1.0            # clip grad norm; prevents rare-class / heavy-tail gradient spikes -> NaN
MAX_CLASS_WEIGHT = 5.0    # cap balanced class weights (singletons would otherwise get ~18x)
EARLY_STOP_PATIENCE = 20
SEED = 42
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# ---------------------------------------------------------------------------
# Model (ported from training/cnn_lstm_attention.py)
# ---------------------------------------------------------------------------


class PositionalEncoding(nn.Module):
    """Standard sinusoidal positional encoding (Vaswani et al.), added before self-attention.

    Attention has no inherent notion of sequence order; without this, the attn blocks are
    permutation-invariant over time and only see order indirectly via the conv backbone's
    local receptive field. Fixed (non-learned) so it needs no extra parameters and
    extrapolates to any T at inference.
    """

    def __init__(self, dim: int, max_len: int = 512):
        super().__init__()
        pos = torch.arange(max_len).unsqueeze(1).float()
        div = torch.exp(torch.arange(0, dim, 2).float() * (-np.log(10000.0) / dim))
        pe = torch.zeros(max_len, dim)
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe.unsqueeze(0))  # (1, max_len, dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.pe[:, : x.size(1)]


def make_transformer_encoder(
    dim: int = 384, heads: int = 8, num_layers: int = 3,
    ff_dim: int | None = None, dropout: float = 0.1,
) -> nn.TransformerEncoder:
    """Standard post-norm Transformer encoder: each of `num_layers` layers is a full
    self-attn sublayer (+residual+LayerNorm) followed by a GELU FFN sublayer
    (+residual+LayerNorm) -- not a bare attention block. `ff_dim` defaults to 4x `dim`,
    the standard Transformer ratio. Input/return: (B, T, dim).
    """
    ff_dim = ff_dim or dim * 4
    layer = nn.TransformerEncoderLayer(
        d_model=dim, nhead=heads, dim_feedforward=ff_dim,
        dropout=dropout, activation="gelu", batch_first=True, norm_first=False,
    )
    return nn.TransformerEncoder(layer, num_layers=num_layers)


class SkatingActionClassifier(nn.Module):
    """(B, T, F) sequence of skeleton features -> (B, num_classes) class logits."""

    def __init__(self, in_features: int, num_classes: int, num_layers: int = 3):
        super().__init__()
        # Conv backbone (expects channels-first (B, F, T))
        self.stem = nn.Conv1d(in_features, 128, 3, padding="same")
        self.stem_bn = nn.BatchNorm1d(128)
        self.cb1 = ConvBlock(128, 192, 3)
        self.cb2 = ConvBlock(192, 256, 3)
        self.cb3 = ConvBlock(256, 384, 5)
        # Positional encoding + Transformer encoder (self-attn + FFN per layer, (B, T, 384))
        self.pos_enc = PositionalEncoding(384)
        self.encoder = make_transformer_encoder(dim=384, heads=8, num_layers=num_layers)
        # Deep classification head (temporal mean-pool -> 384). The reference concatenated two
        # identical mean-pools (768); a single mean-pool carries the same information.
        self.head = nn.Sequential(
            DenseBlock(384, 1024, 0.5),
            DenseBlock(1024, 512, 0.4),
            DenseBlock(512, 256, 0.3),
            nn.LayerNorm(256),
        )
        self.out = nn.Linear(256, num_classes)  # logits; softmax handled by CrossEntropyLoss

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.transpose(1, 2)                      # (B, F, T)
        x = F.relu(self.stem(x))
        x = self.stem_bn(x)
        x = self.cb3(self.cb2(self.cb1(x)))        # (B, 384, T)
        x = x.transpose(1, 2)                      # (B, T, 384)
        x = self.pos_enc(x)
        x = self.encoder(x)
        x = x.mean(dim=1)                          # temporal average pool -> (B, 384)
        return self.out(self.head(x))


def train(data_dir: Path, coarse: bool = False, num_layers: int = 3) -> dict:
    torch.manual_seed(SEED)
    np.random.seed(SEED)

    saved_n, taxonomy = resolve_taxonomy(data_dir, coarse)  # cross-checked against metadata.json
    num_classes = len(taxonomy)
    Xtr, ytr = load_split(data_dir, "train")
    Xva, yva = load_split(data_dir, "val")
    Xte, yte = load_split(data_dir, "test")
    # Saved labels are indices into whichever taxonomy resolve_taxonomy resolved -> validate in
    # that space first, then coarsen (fine-taxonomy data only). This keeps the whole chain
    # consistent: saved idx -> (FINE_TO_COARSE_IDX if coarsening) -> softmax unit.
    assert_label_consistency(saved_n, ytr, yva, yte)
    if coarse:
        def coarsen(y):
            return np.array([labels_mod.FINE_TO_COARSE_IDX[int(v)] for v in y], dtype=np.int64)
        ytr, yva, yte = coarsen(ytr), coarsen(yva), coarsen(yte)
        assert_label_consistency(num_classes, ytr, yva, yte)
    in_features = Xtr.shape[-1]
    if coarse:
        label_space_name = "COARSE action-level"
    elif taxonomy is labels_mod.FS_JUMP3D_TAXONOMY:
        label_space_name = "FS_JUMP3D (7-class jump type)"
    elif taxonomy is labels_mod.FS_JUMP3D_SINGLES_TAXONOMY:
        label_space_name = "FS_JUMP3D_SINGLES (Comb excluded)"
    else:
        label_space_name = "FINE"
    print(f"label space: {label_space_name} | {num_classes} classes")

    # Per-feature standardization from TRAIN stats only (applied to all splits + stored for
    # inference). The velocity block is heavy-tailed (|v| up to ~960 vs coords ~1); feeding it raw
    # blows up activations -> NaN. z-scoring brings every feature onto a comparable scale.
    mu = Xtr.mean(axis=(0, 1), keepdims=True)
    sd = Xtr.std(axis=(0, 1), keepdims=True) + 1e-6
    Xtr = (Xtr - mu) / sd
    Xva = (Xva - mu) / sd
    Xte = (Xte - mu) / sd

    print(f"device={DEVICE} | num_classes={num_classes} (== len(taxonomy)) | in_features={in_features}")
    print(f"shapes: train={Xtr.shape} val={Xva.shape} test={Xte.shape}")
    print(f"standardized features: train |max|={np.abs(Xtr).max():.1f} (was heavy-tailed pre-scaling)")
    print(f"train classes present: {sorted(set(ytr.tolist()))}")

    # class-balanced loss weights (train distribution), capped so singleton classes don't get
    # ~18x weight and destabilize training. "proper" imbalance handling without the spikes.
    present = np.unique(ytr)
    cw = np.clip(compute_class_weight(class_weight="balanced", classes=present, y=ytr),
                 None, MAX_CLASS_WEIGHT)
    weight = torch.ones(num_classes)
    for c, w in zip(present, cw):
        weight[int(c)] = float(w)
    criterion = nn.CrossEntropyLoss(weight=weight.to(DEVICE))

    model = SkatingActionClassifier(in_features, num_classes, num_layers=num_layers).to(DEVICE)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"model params: {n_params/1e6:.2f}M | encoder layers: {num_layers}")
    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)

    tr_loader = DataLoader(
        TensorDataset(torch.from_numpy(Xtr), torch.from_numpy(ytr)),
        batch_size=BATCH_SIZE, shuffle=True, drop_last=False,
    )
    va_loader = DataLoader(
        TensorDataset(torch.from_numpy(Xva), torch.from_numpy(yva)),
        batch_size=BATCH_SIZE, shuffle=False,
    )

    best_val_f1, best_state, since_improved = -1.0, None, 0
    for epoch in range(1, EPOCHS + 1):
        tr_loss, tr_acc = run_epoch(model, tr_loader, criterion, optimizer)
        va_loss, va_acc = run_epoch(model, va_loader, criterion)
        va_f1 = f1_score(yva, predict(model, Xva), labels=list(range(num_classes)),
                         average="macro", zero_division=0)
        if va_f1 > best_val_f1:
            best_val_f1, since_improved = va_f1, 0
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
        else:
            since_improved += 1
        if epoch % 5 == 0 or epoch == 1:
            print(f"epoch {epoch:3d} | train loss {tr_loss:.3f} acc {tr_acc:.3f} "
                  f"| val loss {va_loss:.3f} acc {va_acc:.3f} f1(macro) {va_f1:.3f}")
        if since_improved >= EARLY_STOP_PATIENCE:
            print(f"early stop at epoch {epoch} (no val-F1 improvement for {EARLY_STOP_PATIENCE})")
            break

    if best_state is not None:
        model.load_state_dict(best_state)
    print(f"\nrestored best model (val macro-F1 = {best_val_f1:.3f})")

    metrics = report_metrics(yte, predict(model, Xte), taxonomy)
    layer_suffix = "" if num_layers == 3 else f"_{num_layers}L"
    ckpt_name = f"model_transformer{'_coarse' if coarse else ''}{layer_suffix}.pt"
    torch.save({"state_dict": model.state_dict(), "num_classes": num_classes,
                "in_features": in_features, "taxonomy": taxonomy, "coarse": coarse,
                "num_layers": num_layers,
                "feature_mean": mu, "feature_std": sd},  # standardizer for inference consistency
               data_dir / ckpt_name)
    print(f"\nsaved model -> {data_dir / ckpt_name}")
    return metrics


def smoke() -> None:
    """Validate the full plumbing on synthetic data shaped exactly like the pipeline output."""
    print("SMOKE: synthetic (N,64,94) data, labels in TAXONOMY index space")
    num_classes = len(labels_mod.TAXONOMY)
    rng = np.random.default_rng(0)
    X = rng.standard_normal((40, 64, 94)).astype(np.float32)
    y = rng.integers(0, num_classes, size=40).astype(np.int64)
    assert_label_consistency(num_classes, y)
    model = SkatingActionClassifier(94, num_classes).to(DEVICE)
    logits = model(torch.from_numpy(X[:4]).to(DEVICE))
    assert logits.shape == (4, num_classes), logits.shape
    loss = nn.CrossEntropyLoss()(logits, torch.from_numpy(y[:4]).to(DEVICE))
    loss.backward()
    print(f"forward OK: logits {tuple(logits.shape)} | loss {loss.item():.3f} | backward OK")
    print(f"num_classes={num_classes} matches softmax units; argmax range in [0,{num_classes-1}]")
    print("SMOKE PASSED")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-dir", type=Path, default=Path(DEFAULT_DATA_DIR))
    parser.add_argument("--coarse", action="store_true",
                        help="collapse jump rotations into action-level classes (11 instead of 28)")
    parser.add_argument("--layers", type=int, default=3,
                        help="number of Transformer encoder layers (default 3)")
    parser.add_argument("--smoke", action="store_true", help="self-test on synthetic data")
    args = parser.parse_args()
    if args.smoke:
        smoke()
        return 0
    if not (args.data_dir / "train_features.pkl").exists():
        raise SystemExit(f"No processed data at {args.data_dir}. Run the pipeline first.")
    train(args.data_dir, coarse=args.coarse, num_layers=args.layers)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())