From 85d796b84fb2ca1287e7859e716d0f058fd6730f Mon Sep 17 00:00:00 2001 From: LeoMassai <126597742+LeoMassai@users.noreply.github.com> Date: Mon, 9 Feb 2026 17:28:31 +0100 Subject: [PATCH] Add minimal LRA ListOps benchmark runner --- README.md | 3 + scripts/run_lra_listops.py | 191 +++++++++++++++++++++++++++++++++++++ 2 files changed, 194 insertions(+) create mode 100644 scripts/run_lra_listops.py diff --git a/README.md b/README.md index 7b4ff8e..2c9b821 100644 --- a/README.md +++ b/README.md @@ -89,6 +89,9 @@ The SSM model is implemented by the class DeepSSM, which takes a number of input For a complete, runnable training example on a nonlinear benchmark dataset, see: `Test_files/Tutorial_DeepSSM.py` +For a minimal Long Range Arena (LRA) benchmark script (ListOps) using `DeepSSM`, see: +`scripts/run_lra_listops.py` (requires `pip install datasets`). + ### Tensor shapes and forward outputs - Input tensor shape is `u: (B, L, d_input)` where: diff --git a/scripts/run_lra_listops.py b/scripts/run_lra_listops.py new file mode 100644 index 0000000..9405fb9 --- /dev/null +++ b/scripts/run_lra_listops.py @@ -0,0 +1,191 @@ +"""Minimal LRA ListOps benchmark runner for SSM models. + +Example: + python scripts/run_lra_listops.py --train-steps 200 --eval-steps 50 + +Requires: + pip install datasets +""" + +from __future__ import annotations + +import argparse +import math +from itertools import islice + +import torch +from datasets import load_dataset +from torch import nn +from torch.utils.data import DataLoader +from tqdm import tqdm + +from neural_ssm import DeepSSM + + +class SSMListOpsClassifier(nn.Module): + def __init__( + self, + *, + vocab_size: int, + num_classes: int, + d_model: int, + d_state: int, + n_layers: int, + dropout: float, + ) -> None: + super().__init__() + self.embedding = nn.Embedding(vocab_size, d_model) + self.ssm = DeepSSM( + d_input=d_model, + d_output=d_model, + d_model=d_model, + d_state=d_state, + n_layers=n_layers, + dropout=dropout, + ) + self.classifier = nn.Linear(d_model, num_classes) + + def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + x = self.embedding(input_ids) + x = self.ssm(x) + mask = attention_mask.unsqueeze(-1) + x = (x * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0) + return self.classifier(x) + + +def build_collate_fn(max_length: int, pad_token_id: int) -> callable: + def collate(batch: list[dict[str, object]]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + sequences = [torch.tensor(example["input_ids"][:max_length], dtype=torch.long) for example in batch] + labels = torch.tensor([example["label"] for example in batch], dtype=torch.long) + + lengths = [seq.numel() for seq in sequences] + max_len = max(lengths) + padded = torch.full((len(sequences), max_len), pad_token_id, dtype=torch.long) + attention = torch.zeros((len(sequences), max_len), dtype=torch.float32) + + for idx, seq in enumerate(sequences): + padded[idx, : seq.numel()] = seq + attention[idx, : seq.numel()] = 1.0 + + return padded, attention, labels + + return collate + + +def infer_vocab_size(dataset, max_scan: int) -> int: + feature = dataset.features.get("input_ids") + if feature is not None and hasattr(feature, "feature"): + inner = feature.feature + if getattr(inner, "num_classes", None): + return inner.num_classes + + max_token = 0 + for example in islice(dataset, max_scan): + max_token = max(max_token, max(example["input_ids"])) + return max_token + 1 + + +def infer_num_classes(dataset) -> int: + feature = dataset.features.get("label") + if feature is not None and getattr(feature, "num_classes", None): + return feature.num_classes + labels = dataset["label"] + return len(set(labels)) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Run a minimal LRA ListOps benchmark with DeepSSM.") + parser.add_argument("--batch-size", type=int, default=8) + parser.add_argument("--max-length", type=int, default=2048) + parser.add_argument("--train-steps", type=int, default=200) + parser.add_argument("--eval-steps", type=int, default=50) + parser.add_argument("--lr", type=float, default=3e-4) + parser.add_argument("--d-model", type=int, default=128) + parser.add_argument("--d-state", type=int, default=64) + parser.add_argument("--n-layers", type=int, default=2) + parser.add_argument("--dropout", type=float, default=0.1) + parser.add_argument("--max-vocab-scan", type=int, default=50000) + parser.add_argument("--seed", type=int, default=17) + parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu") + args = parser.parse_args() + + torch.manual_seed(args.seed) + + dataset = load_dataset("lra", "listops") + train_dataset = dataset["train"] + eval_dataset = dataset["validation"] + + vocab_size = infer_vocab_size(train_dataset, args.max_vocab_scan) + num_classes = infer_num_classes(train_dataset) + + model = SSMListOpsClassifier( + vocab_size=vocab_size, + num_classes=num_classes, + d_model=args.d_model, + d_state=args.d_state, + n_layers=args.n_layers, + dropout=args.dropout, + ).to(args.device) + + collate_fn = build_collate_fn(args.max_length, pad_token_id=0) + train_loader = DataLoader( + train_dataset, + batch_size=args.batch_size, + shuffle=True, + collate_fn=collate_fn, + ) + eval_loader = DataLoader( + eval_dataset, + batch_size=args.batch_size, + shuffle=False, + collate_fn=collate_fn, + ) + + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) + loss_fn = nn.CrossEntropyLoss() + + model.train() + train_iter = iter(train_loader) + for step in tqdm(range(1, args.train_steps + 1), desc="train", unit="step"): + try: + input_ids, attention_mask, labels = next(train_iter) + except StopIteration: + train_iter = iter(train_loader) + input_ids, attention_mask, labels = next(train_iter) + + input_ids = input_ids.to(args.device) + attention_mask = attention_mask.to(args.device) + labels = labels.to(args.device) + + logits = model(input_ids, attention_mask) + loss = loss_fn(logits, labels) + + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + + if step % max(1, math.ceil(args.train_steps / 10)) == 0: + tqdm.write(f"step {step}: loss={loss.item():.4f}") + + model.eval() + correct = 0 + total = 0 + with torch.no_grad(): + for step, (input_ids, attention_mask, labels) in enumerate(eval_loader, start=1): + if step > args.eval_steps: + break + input_ids = input_ids.to(args.device) + attention_mask = attention_mask.to(args.device) + labels = labels.to(args.device) + + logits = model(input_ids, attention_mask) + preds = logits.argmax(dim=-1) + correct += (preds == labels).sum().item() + total += labels.numel() + + accuracy = correct / total if total else 0.0 + print(f"eval accuracy (first {total} examples): {accuracy:.4f}") + + +if __name__ == "__main__": + main()