| """Re-generate eval outputs from the MERGED model with safer decoding (longer cap, stop at </svg>, anti-repetition).""" |
| import os, json, time, torch |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
| from huggingface_hub import hf_hub_download, HfApi |
| TOKEN = os.environ["HF_TOKEN"]; REPO = "aixonlab/svg-slm-4b" |
| BS = int(os.environ.get("GEN_BS", 30)); MAXNEW = int(os.environ.get("MAXNEW", 6000)); LIMIT = float(os.environ.get("GEN_SECONDS", 1800)) |
| t0 = time.time() |
| prompts = json.load(open(hf_hub_download("aixonlab/svg-slm-scripts", "eval_prompts.json", token=TOKEN))) |
| tok = AutoTokenizer.from_pretrained(REPO, token=TOKEN); tok.padding_side = "left" |
| model = AutoModelForCausalLM.from_pretrained(REPO, torch_dtype=torch.bfloat16, token=TOKEN, device_map={"": 0}).eval() |
| out = {} |
| for i in range(0, len(prompts), BS): |
| if time.time() - t0 > LIMIT: print("GEN TIME LIMIT", flush=True); break |
| batch = prompts[i:i + BS] |
| texts = [] |
| for p in batch: |
| m = [{"role": "user", "content": p["prompt"]}] |
| try: texts.append(tok.apply_chat_template(m, tokenize=False, add_generation_prompt=True, enable_thinking=False)) |
| except TypeError: texts.append(tok.apply_chat_template(m, tokenize=False, add_generation_prompt=True)) |
| enc = tok(texts, return_tensors="pt", padding=True, add_special_tokens=False).to("cuda") |
| with torch.no_grad(): |
| g = model.generate(**enc, max_new_tokens=MAXNEW, do_sample=True, temperature=0.4, top_p=0.9, |
| repetition_penalty=1.06, stop_strings=["</svg>"], tokenizer=tok, pad_token_id=tok.pad_token_id) |
| for p, seq in zip(batch, g): out[p["id"]] = tok.decode(seq[enc["input_ids"].shape[1]:], skip_special_tokens=True) |
| print("generated", len(out), f"{time.time()-t0:.0f}s", flush=True) |
| json.dump(out, open("/workspace/eval_outputs_v2.json", "w")) |
| HfApi(token=TOKEN).upload_file(path_or_fileobj="/workspace/eval_outputs_v2.json", path_in_repo="eval_outputs_v2.json", repo_id="aixonlab/svg-slm-4b-lora") |
| print("GEN2 DONE", flush=True) |
|
|
