| 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] |
| 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) |
|
|
