| """Generate SVGs for the eval prompts with base + LoRA adapter. Runs on the pod after training.""" | |
| import os, json, time, torch | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| from peft import PeftModel | |
| from huggingface_hub import hf_hub_download, HfApi | |
| TOKEN = os.environ["HF_TOKEN"] | |
| BASE = os.environ.get("BASE", "Qwen/Qwen3.5-4B") | |
| REPO = os.environ.get("OUT_REPO", "aixonlab/svg-slm-4b-lora") | |
| BS = int(os.environ.get("GEN_BS", 20)) | |
| MAXNEW = int(os.environ.get("MAXNEW", 3500)) | |
| GEN_SECONDS = float(os.environ.get("GEN_SECONDS", 1200)) | |
| t0 = time.time() | |
| prompts = json.load(open(hf_hub_download("aixonlab/svg-slm-scripts", "eval_prompts.json", token=TOKEN))) | |
| tok = AutoTokenizer.from_pretrained(BASE, token=TOKEN) | |
| tok.padding_side = "left" | |
| model = AutoModelForCausalLM.from_pretrained(BASE, torch_dtype=torch.bfloat16, token=TOKEN, device_map={"": 0}) | |
| model = PeftModel.from_pretrained(model, "/workspace/adapter").merge_and_unload().eval() | |
| MERGED = os.environ.get("MERGED_REPO", "aixonlab/svg-slm-4b") | |
| try: | |
| api = HfApi(token=TOKEN); api.create_repo(MERGED, exist_ok=True) | |
| model.save_pretrained("/workspace/merged", safe_serialization=True); tok.save_pretrained("/workspace/merged") | |
| api.upload_folder(folder_path="/workspace/merged", repo_id=MERGED, commit_message="merged bf16 weights (base + LoRA)") | |
| print("MERGED PUSHED", flush=True) | |
| except Exception as e: | |
| print("MERGED PUSH FAILED", e, flush=True) | |
| out = {} | |
| for i in range(0, len(prompts), BS): | |
| if time.time() - t0 > GEN_SECONDS: | |
| 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.6, top_p=0.95, | |
| repetition_penalty=1.02, 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.json", "w")) | |
| HfApi(token=TOKEN).upload_file(path_or_fileobj="/workspace/eval_outputs.json", path_in_repo="eval_outputs.json", repo_id=REPO) | |
| print("GEN DONE", flush=True) | |