import os, sys
os.environ.setdefault("ANNULUS_YEAR_OUTPUT","0")   # cap 固化默认
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)
sys.path.insert(0,S)
import cap_test_eval_arrow as CAP        # 用 cap 的 load_model/generate(被 copying=0.96 证过能生成)
_md = os.environ["CKPT"]                 # cap.load_model(nemo) 要含 weights/+context/ 的父目录
if _md.rstrip("/").endswith("/weights"): _md = os.path.dirname(_md.rstrip("/"))
generate, ht = CAP.load_model(_md, "nemo")
P=["The Company's total revenue for the fiscal year",
   "Net income increased primarily due to",
   "The consolidated financial statements",
   "The capital of France is"]
for p in P:
    out = generate(p, max_new_tokens=45)   # cap 的原装 generate(greedy)
    print(f"[cap-gen] {p!r} -> {out!r}", flush=True)
print("CAPGEN_DONE", flush=True)