Test A5 length generalisation

length_gen.py

@@ -0,0 +1,274 @@
+import argparse
+import os
+
+import torch
+
+from data_dir.dataloaders import create_group_dataloaders
+
+
+def build_model(model_name, config, data_dim, label_dim, device):
+    """
+    Build and return a model instance based on 'model_name' & 'config'.
+    """
+
+    if model_name == "mamba":
+        from models.mamba import StackedMamba
+
+        model = StackedMamba(
+            num_blocks=config["num_blocks"],
+            model_dim=config["model_dim"],
+            data_dim=data_dim,
+            label_dim=label_dim,
+            dropout_rate=config.get("dropout_rate", 0.01),
+            use_glu=config.get("use_glu", False),
+            second_embedding=config.get("second_embedding", False),
+        )
+    elif model_name == "s6":
+        from models.s6 import S6
+
+        model = S6(
+            num_blocks=config["num_blocks"],
+            model_dim=config["model_dim"],
+            data_dim=data_dim,
+            label_dim=label_dim,
+            dropout_rate=config.get("dropout_rate", 0.01),
+            use_glu=config.get("use_glu", False),
+            second_embedding=config.get("second_embedding", False),
+            d_state=config.get("d_state", 16),
+            dt_rank=config.get("dt_rank", "auto"),
+        )
+
+    elif model_name == "transformer":
+        from models.transformer import CausalTransformer
+
+        model = CausalTransformer(
+            num_blocks=config["num_blocks"],
+            model_dim=config["model_dim"],
+            data_dim=data_dim,
+            label_dim=label_dim,
+            dropout_rate=config.get("dropout_rate", 0.01),
+            second_embedding=config.get("second_embedding", False),
+        )
+
+    elif model_name in ["deltanet", "gateddeltanet", "deltaproduct", "rwkv6", "rwkv7"]:
+        from models.fla import StackedBlock
+
+        model = StackedBlock(
+            layer_type=model_name,  # "deltanet", "gateddeltanet", "rwkv6", or "rwkv7"
+            num_blocks=config["num_blocks"],
+            model_dim=config["model_dim"],
+            data_dim=data_dim,
+            label_dim=label_dim,
+            rank=config.get("rank", 0),
+            gated=config.get("gated", True),
+            dropout_rate=config.get("dropout_rate", 0.01),
+            use_glu=config.get("use_glu", False),
+            second_embedding=config.get("second_embedding", False),
+        )
+
+    elif model_name == "deltanet2":
+        from models.deltanet2 import StackedBlock as StackedDeltaNet2
+
+        model = StackedDeltaNet2(
+            num_blocks=config["num_blocks"],
+            model_dim=config["model_dim"],
+            data_dim=data_dim,
+            label_dim=label_dim,
+            sigmoid_scale=2,
+            dropout_rate=config.get("dropout_rate", 0.01),
+            use_glu=config.get("use_glu", False),
+            second_embedding=config.get("second_embedding", False),
+        )
+
+    elif model_name == "lcde":
+        from models.slcde import StackedLCDE
+
+        model = StackedLCDE(
+            num_blocks=config["num_blocks"],
+            model_dim=config["model_dim"],
+            data_dim=data_dim,
+            label_dim=label_dim,
+            init_std=config.get("init_std", 1.0),
+            block_size=config.get("block_size", 1),
+            sparsity=config.get("sparsity", 1.0),
+            dropout_rate=config.get("dropout_rate", 0.01),
+            use_glu=config.get("use_glu", False),
+            diagonal=config.get("diagonal", False),
+            fwht=config.get("fwht", False),
+            second_embedding=config.get("second_embedding", False),
+            rank=config.get("rank", 0),
+        )
+
+    elif model_name == "lstm":
+        from models.lstm import LSTM
+
+        model = LSTM(
+            num_blocks=config["num_blocks"],
+            data_dim=data_dim,
+            model_dim=config["model_dim"],
+            label_dim=label_dim,
+            dropout_rate=config.get("dropout_rate", 0.01),
+            second_embedding=config.get("second_embedding", False),
+        )
+
+    elif model_name == "xlstm":
+        from models.xlstm import xLSTM
+
+        model = xLSTM(
+            num_blocks=config["num_blocks"],
+            data_dim=data_dim,
+            model_dim=config["model_dim"],
+            label_dim=label_dim,
+            slstm_at=config.get("slstm_at", [1]),
+            dropout_rate=config.get("dropout_rate", 0.01),
+            second_embedding=config.get("second_embedding", False),
+            context_length=config.get("padding_length", 256),
+        )
+
+    else:
+        raise ValueError(f"Unknown model_name: {model_name}")
+
+    model.to(device)
+    return model
+
+
+def evaluate_model_on_length(
+    model, seq_length, batch_size, num_samples, token_accuracy=False, device="cpu"
+):
+    """
+    Creates a dataloader for 'task' sequences of exactly seq_length,
+    evaluates model on it, and returns the accuracy.
+    """
+    val_dataloader, _, data_dim, label_dim = create_group_dataloaders(
+        group="A5",
+        num_samples=num_samples,
+        batch_size=batch_size,
+        min_length=seq_length,
+        max_length=seq_length,
+        padding_length=seq_length,
+        train_split=1.0,
+        seed=123465,
+    )
+
+    model.eval()
+    correct = 0
+    total = 0
+    if token_accuracy:
+        with torch.no_grad():
+            for X_val, y_val, mask_val in val_dataloader:
+                X_val = X_val.to(device)
+                y_val = y_val.to(device)
+                mask_val = mask_val.to(device)
+
+                outputs_val = model(X_val)  # (batch, seq, label_dim)
+                pred = outputs_val[mask_val].argmax(dim=-1)
+                labels = y_val[mask_val].flatten()
+
+                correct += (pred == labels).sum().item()
+                total += labels.size(0)
+    else:
+        with torch.no_grad():
+            for X_val, y_val, mask_val in val_dataloader:
+                X_val = X_val.to(device)
+                y_val = y_val.to(device)
+                mask_val = mask_val.to(device)
+
+                logits = model(X_val)
+                token_correct = (logits.argmax(-1) == y_val) & mask_val  # (batch, seq)
+                seq_correct = token_correct.all(dim=1)  #    (batch,)
+                correct += seq_correct.sum().item()
+                total += seq_correct.size(0)
+
+    accuracy = correct / total if total > 0 else 0.0
+    return accuracy
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--checkpoint", type=str, required=True, help="Path to the checkpoint .pt file"
+    )
+    parser.add_argument(
+        "--model_name",
+        type=str,
+        required=True,
+        choices=[
+            "xlstm",
+            "s6",
+            "mamba",
+            "lstm",
+            "deltanet",
+            "deltanet2",
+            "gateddeltanet",
+            "deltaproduct",
+            "rwkv6",
+            "rwkv7",
+            "lcde",
+            "transformer",
+        ],
+        help="Which model to load (must match how it was originally trained).",
+    )
+    parser.add_argument(
+        "--lengths",
+        type=str,
+        default="40,56,72,88,104,120,136,152,168,184,200,216,232,248",
+        help="Comma-separated list of sequence lengths to test.",
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=64, help="Batch size for evaluation."
+    )
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        default=2048,
+        help="Number of sequences to sample for each length.",
+    )
+    parser.add_argument(
+        "--strict_load",
+        action="store_true",
+        help="If set, uses strict=True when loading state_dict.",
+    )
+    parser.add_argument(
+        "--token_accuracy",
+        action="store_true",
+        help="If set, evaluates token accuracy instead of sequence accuracy.",
+    )
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    if not os.path.exists(args.checkpoint):
+        raise FileNotFoundError(f"Checkpoint not found: {args.checkpoint}")
+
+    checkpoint = torch.load(args.checkpoint, map_location=device)
+    model_sd = checkpoint["model_state_dict"]
+    config = checkpoint["config"]
+    data_dim = checkpoint["data_dim"]
+    label_dim = checkpoint["label_dim"]
+
+    model = build_model(args.model_name, config, data_dim, label_dim, device)
+
+    if args.strict_load:
+        model.load_state_dict(model_sd, strict=True)
+    else:
+        model.load_state_dict(model_sd, strict=False)
+
+    lengths = [int(x) for x in args.lengths.split(",")]
+    print(f"Evaluating '{args.model_name}' on group='A5' for lengths={lengths}...")
+
+    for L in lengths:
+        acc = evaluate_model_on_length(
+            model,
+            seq_length=L,
+            batch_size=args.batch_size,
+            num_samples=args.num_samples,
+            token_accuracy=args.token_accuracy,
+            device=device,
+        )
+        print(f"Length={L}: Accuracy={acc:.4f}")
+
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()