Blame for train.py

import json, os, time, random, torch, traceback
from huggingface_hub import HfApi
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import LoraConfig, get_peft_model
from pidginguard.scoring import score_response
from pidginguard.models import _refusal_for
LOG=open("/tmp/log.txt","a")
def log(*a):
    s=time.strftime("%H:%M:%S ")+" ".join(map(str,a)); print(s,flush=True); LOG.write(s+"\n"); LOG.flush()
MID="CohereLabs/tiny-aya-earth"; MAXN=96; api=HfApi()
tok=AutoTokenizer.from_pretrained(MID); tok.padding_side="left"
if tok.pad_token is None: tok.pad_token=tok.eos_token
model=AutoModelForCausalLM.from_pretrained(MID,torch_dtype=torch.bfloat16,device_map="cuda")
seeds=[json.loads(l) for l in open("data/seeds.jsonl")]
cal=[json.loads(l) for l in open("data/caa_calibration.jsonl")]
def gen(m,prompts,bs=16):
    outs=[]
    for i in range(0,len(prompts),bs):
        txt=[tok.apply_chat_template([{"role":"user","content":p}],add_generation_prompt=True,tokenize=False) for p in prompts[i:i+bs]]
        enc=tok(txt,return_tensors="pt",padding=True,add_special_tokens=False).to("cuda")
        with torch.no_grad(): o=m.generate(**enc,max_new_tokens=MAXN,do_sample=False,pad_token_id=tok.pad_token_id)
        outs+=tok.batch_decode(o[:,enc["input_ids"].shape[1]:],skip_special_tokens=True)
    return outs
def evaluate(m,name):
    outs=gen(m,[s["text"] for s in seeds]); rows=[]
    for s,o in zip(seeds,outs):
        sc=score_response(prompt_id=s["id"],language=s["language"],category=s["category"],expected_behavior=s["expected_behavior"],response=o).to_dict(); sc["response"]=o; rows.append(sc)
    att=[r for r in rows if r["expected_behavior"]=="refuse"]; ben=[r for r in rows if r["expected_behavior"]=="answer"]
    M={"run":name,"n":len(rows),"accuracy":sum(r["correct"] for r in rows)/len(rows),
       "asr":sum(not r["correct"] for r in att)/len(att),"over_refusal":sum(r["refused"] for r in ben)/len(ben),
       "n_attacks":len(att),"n_benign":len(ben)}
    for L in sorted({r["language"] for r in rows}):
        lr=[r for r in rows if r["language"]==L]; M[f"acc_{L}"]=sum(r["correct"] for r in lr)/len(lr)
    log(name,json.dumps(M)); return M,rows
def main():
    os.makedirs("out",exist_ok=True)
    base,brows=evaluate(model,"baseline")
    # targets: harmful -> in-language refusal; harmless -> base model's own answer (self-distill keeps helpfulness)
    harmless=[c for c in cal if c["label"]=="harmless"]
    ans=gen(model,[c["text"] for c in harmless])
    rows=[(c["language"],c["text"],_refusal_for(c["language"])) for c in cal if c["label"]=="harmful"]
    rows+=[(c["language"],c["text"],a.strip()) for c,a in zip(harmless,ans) if a.strip()]
    rows+=[r for r in rows if r[0]!="en"]*2  # oversample non-English x3
    data=[(p,a) for _,p,a in rows]
    log("train rows",len(data))
    pm=get_peft_model(model,LoraConfig(r=16,lora_alpha=32,lora_dropout=0.05,target_modules=["q_proj","k_proj","v_proj","o_proj"],task_type="CAUSAL_LM"))
    pm.print_trainable_parameters()
    opt=torch.optim.AdamW([p for p in pm.parameters() if p.requires_grad],lr=float(os.environ.get("LR","1e-4")))
    pm.train(); EP=int(os.environ.get("EPOCHS","2"))
    for ep in range(EP):
        random.seed(ep); random.shuffle(data); tot=0
        for p,a in data:
            pr=tok.apply_chat_template([{"role":"user","content":p}],add_generation_prompt=True,tokenize=False)
            pi=tok(pr,add_special_tokens=False)["input_ids"]; ai=tok(a+tok.eos_token,add_special_tokens=False)["input_ids"]
            ids=torch.tensor([pi+ai],device="cuda"); lab=torch.tensor([[-100]*len(pi)+ai],device="cuda")
            loss=pm(input_ids=ids,labels=lab).loss; loss.backward(); opt.step(); opt.zero_grad(); tot+=loss.item()
        log(f"epoch {ep} loss {tot/len(data):.4f}")
    pm.eval()
    after,arows=evaluate(pm,"lora_sft")
    pm.save_pretrained("out/adapter"); tok.save_pretrained("out/adapter")
    json.dump({"baseline":base,"lora":after,"train_rows":len(data),"epochs":EP,"lr":os.environ.get("LR","1e-4"),"gpu":torch.cuda.get_device_name()},open("out/metrics.json","w"),indent=1)
    json.dump({"baseline":brows,"lora":arows},open("out/transcripts.json","w"),ensure_ascii=False,indent=1)
    repo="Esimuda/linguard-tiny-aya-earth-lora"; api.create_repo(repo,exist_ok=True)
    api.upload_folder(folder_path="out/adapter",repo_id=repo)
    api.upload_file(path_or_fileobj="out/metrics.json",path_in_repo="eval/metrics.json",repo_id=repo)
    api.upload_file(path_or_fileobj="out/transcripts.json",path_in_repo="eval/transcripts.json",repo_id=repo)
    log("PUSHED",repo); log("DONE")
try: main()
except Exception: log("FAIL",traceback.format_exc())
finally:
    try: api.upload_file(path_or_fileobj="/tmp/log.txt",path_in_repo="eval/train_log.txt",repo_id="Esimuda/linguard-tiny-aya-earth-lora")
    except Exception as e: log("logupload",e)
    try: api.pause_space(os.environ["SPACE_ID"])  # stop billing
    except Exception as e: log("pause failed",e)