import os, sys, torch
os.environ.setdefault("ANNULUS_YEAR_OUTPUT","0")
os.environ["ANNULUS_ROUTED_EXPERTS"]="32"; os.environ["ANNULUS_SHARED_EXPERTS"]="1"
os.environ["ANNULUS_SHARED_FFN"]="2048"; os.environ.setdefault("ANNULUS_LAYERS","24")
os.environ.setdefault("ANNULUS_TOPK","8"); os.environ["ANNULUS_GROUPED_GEMM"]="0"; os.environ["ANNULUS_GROUP_AUX"]="1"
S="/gpfs/radev/scratch/xu_hua/lq62/annulus_v4"; _CV7=S+"/code_v7"; _REPO=os.path.expanduser("~/Annulus")
for p in [_REPO+"/eval",_REPO+"/nemo/src",_CV7]:
    if os.path.isdir(p):
        if p in sys.path: sys.path.remove(p)
        sys.path.insert(0,p)
import icl_eval_v5 as V
core,tok=V.build_v5_model_and_tokenizer(os.environ["CKPT"],os.environ["TOK"]); core.eval()
@torch.no_grad()
def probe(prompt):
    ids=tok(prompt,add_special_tokens=False)["input_ids"]
    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=core(input_ids=inp,position_ids=pos,attention_mask=m)
    L = o[0] if o.shape[0]==1 else o[:,0]     # [seq,V]
    print(f"\nPROMPT {prompt!r}  (o.shape={tuple(o.shape)})",flush=True)
    d=lambda t: tok.decode([int(t)])
    for i in range(s):
        top1=int(L[i].float().argmax())
        cur=d(ids[i]); nxt=d(ids[i+1]) if i+1<s else "?"
        flag=""
        if top1==ids[i]: flag=" <== top1==CURRENT(off-by-one!)"
        elif i+1<s and top1==ids[i+1]: flag=" (top1==true_next ✓)"
        print(f"  pos{i}: in={cur!r} top1={d(top1)!r} true_next={nxt!r}{flag}",flush=True)
for p in ["The capital of France is","Water is made of hydrogen and"]:
    probe(p)
print("POSPROBE_DONE",flush=True)