75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
|
"""IOL-AI submission script — RUN 1: pipeline smoke test."""
|
||
|
|
import os
|
||
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
||
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||
|
|
|
||
|
|
import re
|
||
|
|
import json
|
||
|
|
|
||
|
|
MODEL_ID = "."
|
||
|
|
|
||
|
|
SYSTEM = (
|
||
|
|
"You solve International Linguistics Olympiad problems. Everything needed is in the "
|
||
|
|
"problem itself. Answer every numbered item. Put each answer on its own line, in the "
|
||
|
|
"same order the items appear, with no numbering and no extra commentary."
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def count_items(query: str) -> int:
|
||
|
|
line_items = re.findall(r"(?m)^\s*[\(\[]?(\d+)[\.\)\]]", query or "")
|
||
|
|
if line_items:
|
||
|
|
return len(line_items)
|
||
|
|
paren_items = re.findall(r"\((\d+)\)", query or "")
|
||
|
|
if paren_items:
|
||
|
|
return len(set(paren_items))
|
||
|
|
return 1
|
||
|
|
|
||
|
|
|
||
|
|
def _clean(line: str) -> str:
|
||
|
|
line = re.sub(r"^\s*[\(\[]?\d+[\.\)\]]?\s*", "", line)
|
||
|
|
return line.strip().strip('"').strip("'").strip()
|
||
|
|
|
||
|
|
|
||
|
|
def parse_answers(text: str, n: int):
|
||
|
|
lines = [_clean(ln) for ln in (text or "").splitlines()]
|
||
|
|
lines = [ln for ln in lines if ln]
|
||
|
|
if len(lines) >= n:
|
||
|
|
return lines[:n]
|
||
|
|
return lines + [""] * (n - len(lines))
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
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("")
|
||
|
|
out_rows = []
|
||
|
|
for _, r in df.iterrows():
|
||
|
|
n = count_items(r["query"])
|
||
|
|
messages = [
|
||
|
|
{"role": "system", "content": SYSTEM},
|
||
|
|
{"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():
|
||
|
|
gen = model.generate(ids, max_new_tokens=512, do_sample=False)
|
||
|
|
text = tok.decode(gen[0][ids.shape[-1]:], skip_special_tokens=True).strip()
|
||
|
|
answers = parse_answers(text, n)
|
||
|
|
out_rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
|
||
|
|
print(f"{len(out_rows)}/{len(df)} done id={r['id']} items={n}", flush=True)
|
||
|
|
|
||
|
|
pd.DataFrame(out_rows).to_csv("submission.csv", index=False)
|
||
|
|
print("wrote submission.csv", flush=True)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|