Spaces:
Runtime error
Runtime error
| 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) | |