File size: 2,422 Bytes
9e357f4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | import os, subprocess, sys, time
T0 = time.time()
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
MODEL_ID = "."
DEADLINE = 26 * 60 # write CSV by 26 min, before the 30-min kill
MAX_NEW = 320 # answers are short; don't burn clock on 512
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "bitsandbytes"], check=True)
print(f"[setup] deps ok at {time.time()-T0:.0f}s", flush=True)
import json
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16)
tok = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID, quantization_config=bnb, device_map="auto"
).eval()
print(f"[model] loaded at {time.time()-T0:.0f}s", flush=True)
df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
rows = []
for i, (_, r) in enumerate(df.iterrows()):
if time.time() - T0 > DEADLINE:
print(f"[deadline] stopping at {i}/{len(df)}", flush=True)
rows.extend({"id": rr["id"], "pred": json.dumps([])}
for _, rr in df.iloc[i:].iterrows())
break
messages = [
{"role": "system", "content":
"You solve International Linguistics Olympiad problems. Answer every numbered "
"item. Put each answer on its own line, in order, with no numbering, no labels, "
"and no extra text. Answer in the language the item asks for."},
{"role": "user", "content": f"{r['context'].strip()}\n\n{r['query'].strip()}"},
]
ids = tok.apply_chat_template(messages, add_generation_prompt=True,
return_tensors="pt").to(model.device)
with torch.no_grad():
out = model.generate(ids, max_new_tokens=MAX_NEW, do_sample=False,
pad_token_id=tok.eos_token_id)
text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
answers = [ln.strip() for ln in text.splitlines() if ln.strip()]
rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
print(f"{i+1}/{len(df)} done at {time.time()-T0:.0f}s", flush=True)
pd.DataFrame(rows).to_csv("submission.csv", index=False)
print(f"wrote submission.csv at {time.time()-T0:.0f}s", flush=True)
|