"""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)