File size: 1,451 Bytes
3b62c8f 42c1993 3b62c8f 42c1993 3b62c8f 42c1993 | 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 | import os
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
MODEL_ID = "."
import json
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
tok = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID, torch_dtype=torch.float16, device_map="auto"
).eval()
df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
rows = []
for _, r in df.iterrows():
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 and no extra text."},
{"role": "user", "content": f"{r['context'].strip()}\n\n{r['query'].strip()}"},
]
enc = tok.apply_chat_template(
messages, add_generation_prompt=True, return_tensors="pt", return_dict=True,
).to(model.device)
with torch.no_grad():
out = model.generate(**enc, max_new_tokens=512, do_sample=False)
text = tok.decode(out[0][enc["input_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"{len(rows)}/{len(df)} done", flush=True)
pd.DataFrame(rows).to_csv("submission.csv", index=False)
print("wrote submission.csv", flush=True)
|