| """cap_test 5-shot evaluation harness (Annulus base-v4, year-aware). |
| |
| Dataset: lfqian/cap_test — 6 tasks (splits), each row = {input, output, task}: |
| copying, uppercase, basic_arithmetic, plural_upper, gerund_upper, math |
| |
| We evaluate OUR year-aware model under BOTH prompt conventions it was trained on: |
| |
| --mode task1 (year-conditioned generation) |
| Sequence starts with a year token [Yt], then a 5-shot demo block, then the |
| query. The model just continues (plain LM). This is exactly how Task1 |
| samples look: [Yt] <text...>. |
| |
| --mode task2 (self-emit-year-then-continue) |
| Mirrors Task2 layout ctx <answer> [Yt] cont : the demo block is the ctx, |
| then we append <answer>, force-feed the correct [Yt], and let the model |
| produce the answer as `cont`. Tests that Task2 wiring didn't break the |
| model's ability to still do the surface task after emitting the year token. |
| |
| Both share the same 5-shot demonstrations and exact-match scoring; the only |
| difference is where the [Yt]/<answer> special tokens sit — so a gap between the |
| two modes localises to the year-token machinery, not the task itself. |
| |
| 5-shot: for each query we sample 5 OTHER (input,output) pairs from the same task |
| as in-context demonstrations (fixed rng per task so runs are comparable). |
| |
| Output: per-task exact-match accuracy + overall, optional --csv for review. |
| |
| Model loading: our ckpt is a NeMo/Megatron distcp MoE with a custom year router, |
| NOT loadable by AutoModelForCausalLM. Plug the real generate() in `load_model()` |
| (NeMo inference or the offline-inference path). A HF AutoModel branch is provided |
| for a converted ckpt / for running a Qwen baseline for comparison. |
| """ |
| import argparse, csv, json, os, random, sys |
|
|
| |
| |
| |
| |
| os.environ.setdefault("ANNULUS_YEAR_OUTPUT", "0") |
|
|
| |
| YEAR_TOKEN_DEFAULT = "[Y2000]" |
| ANSWER_TOKEN = "<answer>" |
| DATE_TOKEN = "<date?>" |
|
|
| TASKS = ["copying", "uppercase", "basic_arithmetic", |
| "plural_upper", "gerund_upper", "math"] |
|
|
|
|
| |
| def load_cap_test(): |
| """Return {task: [ {input, output}, ... ]} for all 6 splits. |
| |
| CAP_TEST_DIR lets us read the parquet from disk: compute nodes have no |
| internet. Verified md5-identical to lfqian/cap_test on the Hub. |
| """ |
| local = os.environ.get("CAP_TEST_DIR") |
| if local: |
| import pyarrow.parquet as pq, glob |
| out = {} |
| for t in TASKS: |
| hits = glob.glob(os.path.join(local, f"{t}*.parquet")) |
| if not hits: |
| raise SystemExit(f"no parquet for task {t} in {local}") |
| rows = pq.read_table(hits[0]).to_pylist() |
| out[t] = [{"input": r["input"], "output": r["output"]} for r in rows] |
| return out |
| from datasets import load_dataset |
| out = {} |
| for t in TASKS: |
| ds = load_dataset("lfqian/cap_test", split=t) |
| out[t] = [{"input": r["input"], "output": r["output"]} for r in ds] |
| return out |
|
|
|
|
| |
| def build_demos(rows, query_idx, k, rng): |
| """Pick k demonstration rows != query_idx.""" |
| pool = [i for i in range(len(rows)) if i != query_idx] |
| idx = rng.sample(pool, min(k, len(pool))) |
| return [rows[i] for i in idx] |
|
|
|
|
| def render_shots(demos, query_in): |
| |
| |
| body = "".join(f"{d['input']} -> {d['output']}\n" for d in demos) |
| return f"{body}{query_in} ->" |
|
|
|
|
| def make_prompt(mode, year_token, demos, query_in): |
| """Return the string fed to the tokenizer. Arrow format. |
| |
| task1: [Yt] {demos} q -> model emits only the result |
| task2: <date?> {demos} q -> <answer> the arrow is IN the input; after |
| <answer> the model self-emits [Yt] then the result. |
| """ |
| shots = render_shots(demos, query_in) |
| if mode == "plain": |
| return shots |
| if mode == "task1": |
| return f"{year_token}{shots}" |
| elif mode == "task2": |
| return f"{DATE_TOKEN}{shots}{ANSWER_TOKEN}" |
| raise ValueError(mode) |
|
|
|
|
| |
| def load_model(model_path, backend, trust_remote_code=False): |
| """Return a callable generate(prompt:str, max_new_tokens:int)->str.""" |
| if backend == "hf": |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=trust_remote_code) |
| model = AutoModelForCausalLM.from_pretrained( |
| model_path, torch_dtype=torch.bfloat16, device_map="auto", |
| trust_remote_code=trust_remote_code) |
| model.eval() |
|
|
| def generate(prompt, max_new_tokens=24): |
| ids = tok(prompt, return_tensors="pt").to(model.device) |
| with torch.no_grad(): |
| out = model.generate(**ids, max_new_tokens=max_new_tokens, |
| do_sample=False, |
| pad_token_id=tok.eos_token_id) |
| new = out[0][ids["input_ids"].shape[1]:] |
| |
| return tok.decode(new, skip_special_tokens=False) |
| return generate, tok |
|
|
| if backend == "nemo": |
| |
| |
| import types, glob, torch |
| for _v in ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN"): |
| os.environ.pop(_v, None) |
|
|
| _m = "megatron.core.dist_checkpointing.strategies.tensorstore" |
| try: |
| __import__(_m) |
| except Exception: |
| st = types.ModuleType(_m) |
|
|
| class _T: |
| def __init__(s, *a, **k): pass |
| st.TensorStoreLoadShardedStrategy = _T |
| st.TensorStoreSaveShardedStrategy = _T |
| sys.modules[_m] = st |
| import megatron.core.dist_checkpointing.strategies as _s |
| _s.tensorstore = st |
|
|
| import torch.distributed as dist |
| for k, v in dict(MASTER_ADDR="localhost", |
| MASTER_PORT=os.environ.get("MASTER_PORT", "29761"), |
| RANK="0", WORLD_SIZE="1", LOCAL_RANK="0").items(): |
| os.environ.setdefault(k, v) |
| torch.cuda.set_device(0) |
| if not dist.is_initialized(): |
| dist.init_process_group("nccl", rank=0, world_size=1) |
| from megatron.core import parallel_state as ps |
| if not ps.model_parallel_is_initialized(): |
| ps.initialize_model_parallel(1, 1) |
| from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed |
| model_parallel_cuda_manual_seed(1234) |
|
|
| WEIGHTS = os.path.join(model_path, "weights") |
| CTX = os.path.join(model_path, "context") |
| from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer |
| _td = [d for d in sorted(glob.glob(os.path.join(CTX, "*"))) if os.path.isdir(d) |
| and os.path.exists(os.path.join(d, "tokenizer_config.json"))] |
| if not _td: |
| raise SystemExit(f"no tokenizer dir under {CTX}") |
| tok = AutoTokenizer(pretrained_model_name=_td[0]); ht = tok.tokenizer |
| print(f"[nemo] tokenizer {_td[0]}", flush=True) |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| import annulus.model.year_routing as _yr |
| import annulus.model.year_embedding as _ye |
| if _yr.YEAR_TOKEN_BASE_ID is not None: |
| raise SystemExit( |
| f"[baseline] annulus.year_routing.YEAR_TOKEN_BASE_ID is hardcoded to " |
| f"{_yr.YEAR_TOKEN_BASE_ID} -- that is the PRE-v4 package. Point " |
| f"PYTHONPATH at the tf4-launch-fix code (code_tf4/).") |
| _y2000 = _yr.resolve_year_token_base_id(ht) |
| _ye.verify_year_vocab_layout(ht) |
| print(f"[nemo] year ids resolved from tokenizer: [Y2000]={_y2000}", flush=True) |
|
|
| |
| |
| |
| |
| |
| cfg = None |
| try: |
| from nemo.lightning import io as _nlio |
| cfg = _nlio.load_context(CTX, subpath="model.config") |
| print("[nemo] CONFIG: from context/io.json (AnnulusGPTConfig)", flush=True) |
| except Exception as _e: |
| print(f"[nemo] io.load_context failed: {repr(_e)[:200]}", flush=True) |
| from annulus.model.base_config import make_base_config |
| cfg = make_base_config() |
| print("[nemo] CONFIG: make_base_config()", flush=True) |
|
|
| import yaml as _yaml |
| _mp = os.path.join(CTX, "model.yaml") |
| if not os.path.exists(_mp): |
| raise SystemExit(f"no model.yaml in {CTX}; arch cannot be trusted") |
| _y = (_yaml.safe_load(open(_mp)) or {}).get("config", {}) or {} |
| _chg = [] |
| for _k, _v in _y.items(): |
| if _k.startswith("_") or isinstance(_v, dict) or _v is None: continue |
| if isinstance(_v, (list, tuple)) and any(isinstance(_x, (dict, list)) for _x in _v): |
| continue |
| if hasattr(cfg, _k) and getattr(cfg, _k) != _v: |
| try: |
| _chg.append((_k, getattr(cfg, _k), _v)); setattr(cfg, _k, _v) |
| except Exception: pass |
| print(f"[nemo] {len(_chg)} config overrides from model.yaml " |
| f"(0 expected if the baseline matches the ckpt)", flush=True) |
| for _k, _o, _n in _chg: |
| print(f" {_k}: {_o} -> {_n}", flush=True) |
| for _a in ("num_layers", "num_moe_experts", "moe_router_topk", |
| "moe_shared_expert_intermediate_size", "hidden_size"): |
| print(f" cfg.{_a} = {getattr(cfg, _a, 'n/a')}", flush=True) |
| try: |
| from megatron.core.transformer.enums import AttnBackend |
| cfg.attention_backend = AttnBackend.unfused |
| except Exception: pass |
|
|
| |
| |
| |
| |
| |
| |
| from annulus.model.year_aware_model import wire_year_hooks |
| mcore = cfg.configure_model(tok) |
| _yc = wire_year_hooks(mcore, tok, register_on_model=True, verbose=True) |
| _routers = [m for _, m in mcore.named_modules() |
| if m.__class__.__name__ == "AnnulusYearRouter"] |
| if not _routers: |
| raise SystemExit("[wiring] no AnnulusYearRouter after wire_year_hooks") |
| _e = [n for n, _ in mcore.named_parameters() if n.endswith("year_context.e")] |
| if len(_e) != 1: |
| raise SystemExit(f"[wiring] expected 1 year_context.e, found {len(_e)}") |
| print(f"[wiring] AnnulusYearRouter x{len(_routers)}, year_context.e ✓", flush=True) |
| mcore = mcore.cuda().eval() |
| from megatron.core import dist_checkpointing |
| sd = mcore.sharded_state_dict(prefix="module.") |
| loaded = dist_checkpointing.load(sd, WEIGHTS) |
| strip = {(k[7:] if k.startswith("module.") else k): v for k, v in loaded.items()} |
| miss = mcore.load_state_dict(strip, strict=False) |
| real = [m for m in miss.missing_keys if "optim" not in m] |
| print(f"[nemo] loaded {WEIGHTS} missing_nonopt={len(real)}", flush=True) |
| if real[:5]: |
| print(f" e.g. {real[:5]}", flush=True) |
|
|
| EOD = ht.convert_tokens_to_ids("<|endoftext|>") |
|
|
| @torch.no_grad() |
| def generate(prompt, max_new_tokens=24): |
| |
| |
| ids = ht(prompt, add_special_tokens=False)["input_ids"] |
| out = [] |
| for _ in range(max_new_tokens): |
| s = len(ids) |
| inp = torch.tensor([ids], device="cuda") |
| pos = torch.arange(s, device="cuda")[None] |
| m = torch.triu(torch.ones(s, s, dtype=torch.bool, device="cuda"), 1)[None, None] |
| o = mcore(input_ids=inp, position_ids=pos, attention_mask=m) |
| v = (o[0, -1] if o.shape[0] == 1 else o[-1, 0]) |
| nx = int(v.float().argmax()) |
| ids.append(nx); out.append(nx) |
| if nx == EOD: break |
| |
| return ht.decode(out, skip_special_tokens=False) |
|
|
| return generate, ht |
| raise ValueError(backend) |
|
|
|
|
| |
| import re |
| _YEAR_RE = re.compile(r"^\s*(\[Y[^\]]*\]|<beyond_horizon>|<far_future>)\s*") |
| _SPECIAL_RE = re.compile(r"\[Y[^\]]*\]|<answer>|<future>|<far_future>|<beyond_horizon>") |
|
|
|
|
| def split_year(text): |
| """For task2: pull the model-emitted leading year token off the front. |
| Returns (year_token_or_None, remaining_text).""" |
| m = _YEAR_RE.match(text) |
| if m: |
| return m.group(1), text[m.end():] |
| return None, text |
|
|
|
|
| def strip_special(text): |
| """Remove year/special tokens so a bare gold substring can't collide with |
| the digits inside an emitted year token (e.g. '8' inside '[Y1998]').""" |
| return _SPECIAL_RE.sub(" ", text) |
|
|
|
|
| def first_line(text): |
| return text.strip().split("\n")[0].strip() |
|
|
|
|
| def run(args): |
| rng = random.Random(args.seed) |
| data = load_cap_test() |
| generate, _ = load_model(args.model, args.backend, args.trust_remote_code) |
|
|
| rows_out, totals = [], {} |
| |
| def _norm(s): return str(s).split("\n")[0].strip().lower() |
| for task in TASKS: |
| rows = data[task] |
| n_len = 0; n_str = 0 |
| n = min(args.max_per_task, len(rows)) if args.max_per_task else len(rows) |
| task_rng = random.Random(args.seed + hash(task) % 1000) |
| for qi in range(n): |
| q = rows[qi] |
| demos = build_demos(rows, qi, args.shots, task_rng) |
| prompt = make_prompt(args.mode, args.year_token, demos, q["input"]) |
| gold = q["output"].strip() |
| emitted_year = None |
| if args.mode == "task2": |
| raw = generate(prompt, max_new_tokens=args.task2_window) |
| emitted_year, _ = split_year(raw) |
| window = strip_special(raw) |
| pred = window.strip() |
| ok_str = (gold in window); ok_len = (_norm(gold) in window.lower()) |
| else: |
| raw = generate(prompt) |
| pred = first_line(raw) |
| ok_str = (pred == gold) |
| ok_len = (_norm(gold) in raw.strip().lower()) |
| n_str += ok_str; n_len += ok_len |
| if qi < 10: |
| print(f"[SAMPLE] {task:12s} input={q['input']!r} gold={gold!r} pred={pred!r} " |
| f"raw={raw.strip()[:50]!r} strict={int(ok_str)} lenient={int(ok_len)}", flush=True) |
| rows_out.append({"task": task, "mode": args.mode, "input": q["input"], |
| "gold": gold, "pred": pred, "raw": raw.strip()[:80], |
| "lenient": int(ok_len), "strict": int(ok_str), |
| "emitted_year": emitted_year or ""}) |
| totals[task] = (n_len, n_str, n) |
| print(f"[{args.mode}] {task:16s} lenient={n_len/n:.3f} strict={n_str/n:.3f} (n={n})", flush=True) |
|
|
| tl = sum(l for l, _, _ in totals.values()); ts = sum(s for _, s, _ in totals.values()) |
| tn = sum(nn for _, _, nn in totals.values()) |
| print(f"[{args.mode}] {'OVERALL':16s} lenient={tl/tn:.3f} strict={ts/tn:.3f} (n={tn})", flush=True) |
|
|
| if args.csv: |
| with open(args.csv, "w", newline="") as f: |
| w = csv.DictWriter(f, fieldnames=list(rows_out[0].keys())) |
| w.writeheader(); w.writerows(rows_out) |
| print("wrote", args.csv) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("model", help="ckpt path / HF repo (or converted annulus-base-v4)") |
| ap.add_argument("--backend", choices=["hf", "nemo"], default="nemo") |
| ap.add_argument("--mode", choices=["plain", "task1", "task2"], default="task1") |
| ap.add_argument("--shots", type=int, default=5) |
| ap.add_argument("--task2-window", dest="task2_window", type=int, default=8, |
| help="task2: match gold within the first N generated tokens") |
| ap.add_argument("--year-token", dest="year_token", default=YEAR_TOKEN_DEFAULT) |
| ap.add_argument("--max-per-task", dest="max_per_task", type=int, default=0) |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--trust-remote-code", action="store_true") |
| ap.add_argument("--csv") |
| args = ap.parse_args() |
| run(args) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
