Upload gen2.py with huggingface_hub

gen2.py

@@ -0,0 +1,28 @@
+"""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)