scripts/claim5_smoke.py

"""Claim 5 (architecture/params) + end-to-end smoke test on GPU.

Loads the released APT weights, prints the AR model config and parameter
counts, and reconstructs the bundled 1CRN.pdb to confirm the encode/decode
pipeline works before scaling to full test sets.
"""
from __future__ import annotations

import json
import os
from pathlib import Path

import torch

from apt.models import APTLanguageModel, APTTokenizer
from apt.utils import kabsch_rmsd

CKPT = Path(os.environ.get("CKPT_DIR", "/tmp/apt_ckpts"))
TOK = CKPT / "tokenizer128.pt"
LM = CKPT / "lm128.pt"


def load_ca_coords(pdb_path: Path) -> torch.Tensor:
    coords = []
    with pdb_path.open() as handle:
        for line in handle:
            if not line.startswith("ATOM"):
                continue
            if line[12:16].strip() != "CA":
                continue
            if line[16:17] not in (" ", "A"):
                continue
            coords.append([float(line[30:38]), float(line[38:46]), float(line[46:54])])
    x = torch.tensor(coords, dtype=torch.float32)
    x = x - x.mean(dim=0, keepdim=True)
    return x / 10.0


def main() -> None:
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"device={dev} torch={torch.__version__}")

    tokenizer = APTTokenizer.from_pretrained(TOK).to(dev).eval()
    model = APTLanguageModel.from_pretrained(LM, TOK).to(dev).eval()

    cfg = model.cfg
    lm_params = model.num_params()
    tok_params = tokenizer.num_params()

    report = {
        "lm_config": {
            "n_layers": cfg.n_layers,
            "n_channels": cfg.n_channels,
            "n_heads": cfg.n_heads,
            "block_size": cfg.block_size,
            "vocab_size": model.vocab_size,
        },
        "lm_params": int(lm_params),
        "lm_params_M": round(lm_params / 1e6, 3),
        "tokenizer_params": int(tok_params),
        "tokenizer_params_M": round(tok_params / 1e6, 3),
        "tokenizer_cfg": {
            "n_tokens": tokenizer.cfg.n_tokens,
            "levels": list(tokenizer.cfg.levels),
        },
    }

    # Claim 5 checks: 20 layers, 1024 channels, 252.9M params
    report["claim5_layers_ok"] = cfg.n_layers == 20
    report["claim5_channels_ok"] = cfg.n_channels == 1024
    report["claim5_params_252_9M_ok"] = abs(lm_params - 252_893_184) < 100_000

    # Smoke: reconstruct bundled example
    x = load_ca_coords(Path("examples/1CRN.pdb")).unsqueeze(0).to(dev)
    max_toks = tokenizer.cfg.n_tokens
    with torch.no_grad():
        _, _, idx_BL = tokenizer.encode(x)
        idx_BL = idx_BL[:, :max_toks]
        recon = tokenizer.decode(idx_BL)
        rmsd = kabsch_rmsd(recon.cpu(), x.cpu()) * 10
    report["smoke_1CRN_residues"] = int(x.shape[1])
    report["smoke_1CRN_recon_rmsd_A"] = round(float(rmsd.item()), 4)

    print("RESULT_JSON " + json.dumps(report))
    out = Path(os.environ.get("OUT_DIR", "/tmp/out"))
    out.mkdir(parents=True, exist_ok=True)
    (out / "claim5_smoke.json").write_text(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()