"""Self-contained ArithMark-2 verification for Nexus-Erebus models.

Reproduces the reported ArithMark-2 score. No local files needed beyond this repo.
The model ships a digit-atomic, least-significant-digit-first tokenizer, so it must
be loaded with trust_remote_code=True. Text goes in and comes out in normal order;
the digit reversal happens inside the tokenizer.

    pip install torch transformers datasets
    python benchmark_nexus_arithmark.py            # uses this repo
    python benchmark_nexus_arithmark.py <model_id>
"""
import sys, ast, torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = sys.argv[1] if len(sys.argv) > 1 else "."
dev = "cuda" if torch.cuda.is_available() else "cpu"

tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to(dev).eval()
ds = load_dataset("AxiomicLabs/ArithMark-2.0", split="train")


@torch.no_grad()
def avg_logprob(ctx: str, ending: str) -> float:
    """Mean log-prob of `ending` conditioned on `ctx` (the leaderboard's scoring)."""
    ctx_ids = tok(ctx, return_tensors="pt").input_ids.to(dev)
    full_ids = tok(ctx + ending, return_tensors="pt").input_ids.to(dev)
    if full_ids.shape[1] <= ctx_ids.shape[1]:
        return -1e9
    logits = model(full_ids).logits[:, :-1, :]
    logp = torch.log_softmax(logits, dim=-1)
    tgt = full_ids[:, 1:]
    sel = logp.gather(2, tgt.unsqueeze(-1)).squeeze(-1)[:, ctx_ids.shape[1] - 1:]
    return sel.mean().item()


correct = 0
for i, e in enumerate(ds):
    endings = e["endings"] if isinstance(e["endings"], list) else ast.literal_eval(e["endings"])
    scores = [avg_logprob(e["ctx"], end) for end in endings]
    if max(range(len(scores)), key=lambda j: scores[j]) == int(e["label"]):
        correct += 1
    if (i + 1) % 500 == 0:
        print(f"  {i+1}/{len(ds)}  running acc: {correct/(i+1):.4f}", flush=True)

print(f"\nArithMark-2 accuracy for {MODEL}: {correct/len(ds)*100:.2f}%  ({correct}/{len(ds)})")