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()
EOS=tok.eos_token_id
@torch.no_grad()
def gen(prompt,k=50,rep=1.8,top_k=40,top_p=0.95,temp=0.8,nrn=3):
    ids=[151830]+tok(prompt,add_special_tokens=False)["input_ids"]  # prepend [Y2000]=151830 (v11 训练格式)
    start=len(ids)
    for _ in range(k):
        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)
        lg=(o[0,-1] if o.shape[0]==1 else o[-1,0]).float()
        for t in set(ids): lg[t]=lg[t]/rep if lg[t]>0 else lg[t]*rep   # rep penalty
        if len(ids)>=nrn-1:                                            # no_repeat_ngram=3
            pref=tuple(ids[-(nrn-1):])
            for i in range(len(ids)-nrn+1):
                if tuple(ids[i:i+nrn-1])==pref: lg[ids[i+nrn-1]]=-1e9
        lg=lg/temp
        v,i=lg.topk(min(top_k,lg.numel()))                            # top-k
        sp=torch.softmax(v,-1); cum=sp.cumsum(-1); keep=cum<=top_p; keep[0]=True  # top-p
        v=v[keep]; i=i[keep]; pr=torch.softmax(v,-1); nxt=int(i[torch.multinomial(pr,1)])
        ids.append(nxt)
        if EOS is not None and nxt==EOS: break
    return tok.decode(ids[start:])
FIN=["The Company's total revenue for the fiscal year","Net income increased primarily due to",
     "The consolidated financial statements","Our principal risk factors include"]
GEN=["The capital of France is","Water is made of hydrogen and"]
print("== FINANCE (rep1.8/topk40/topp0.95/temp0.8/no_repeat_ngram3) ==",flush=True)
for p in FIN: print(f"[fin] {p!r} -> {gen(p)!r}",flush=True)
print("== GENERAL ==",flush=True)
for p in GEN: print(f"[gen] {p!r} -> {gen(p)!r}",flush=True)
print("TALK2_DONE",flush=True)