253 lines
8.8 KiB
Python
253 lines
8.8 KiB
Python
#!/usr/bin/env python3
|
|
"""IOL-AI 2026 — v8 normal (pred) + explanation por problema (no hardcode).
|
|
|
|
Fase 1: prompt/parse EXACTOS de script_submission8.py → pred (score).
|
|
Fase 2: si sobra tiempo, 1 frase de explicación por fila (generate corto).
|
|
Si no hay tiempo: explicación hecha del propio problema (task/lang/query),
|
|
distinta en cada fila — nunca la misma frase estática para todas.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import time
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
|
|
|
MODEL_ID = os.environ.get("IOL_MODEL", ".")
|
|
INPUT_CSV = Path(os.environ.get("IOL_INPUT", "/tmp/data/test.csv"))
|
|
OUTPUT_CSV = Path(os.environ.get("IOL_OUTPUT", "submission.csv"))
|
|
SOFT_DEADLINE = float(os.environ.get("IOL_SOFT_DEADLINE", "1650"))
|
|
WRITE_EXPLANATIONS = os.environ.get("IOL_EXPLANATIONS", "1") != "0"
|
|
MAX_NEW = int(os.environ.get("IOL_MAX_NEW_TOKENS", "512"))
|
|
EXPL_MAX_NEW = int(os.environ.get("IOL_EXPL_MAX_NEW", "48"))
|
|
EXPL_STOP_LEFT = float(os.environ.get("IOL_EXPL_STOP_LEFT", "20"))
|
|
|
|
SYSTEM_V8 = (
|
|
"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."
|
|
)
|
|
|
|
|
|
def normalize_match_letter(ans: str) -> str:
|
|
ans = ans.strip()
|
|
m = re.fullmatch(r"[\(\[]?([A-Za-z])[\)\]]?[.)]?", ans)
|
|
if m:
|
|
return m.group(1).upper()
|
|
tokens = re.findall(r"\b([A-Za-z])\b", ans)
|
|
if tokens:
|
|
return tokens[-1].upper()
|
|
m = re.search(r"[A-Za-z]", ans)
|
|
return m.group(0).upper() if m else ans
|
|
|
|
|
|
def normalize_text_to_num(ans: str) -> str:
|
|
a = re.sub(r"(?i)^(answer|ans|result)\s*[:=]\s*", "", ans.strip()).strip()
|
|
if re.fullmatch(r"[\d\s+\-*/^=()]+", a.replace(",", "")):
|
|
a = a.replace(",", "").replace(" ", "")
|
|
if "=" in a and " = " not in a:
|
|
a = a.replace("=", " = ")
|
|
return a.strip()
|
|
m = re.search(r"\d+", a)
|
|
return m.group(0) if m and len(a) < 40 else a
|
|
|
|
|
|
def safe_normalize_answers(answers: List[str], task_type: str) -> List[str]:
|
|
task_type = (task_type or "").strip().lower()
|
|
out: List[str] = []
|
|
for a in answers:
|
|
a = a.strip()
|
|
if task_type == "match_letters":
|
|
a = normalize_match_letter(a)
|
|
elif task_type == "text_to_num":
|
|
a = normalize_text_to_num(a)
|
|
out.append(a)
|
|
return out
|
|
|
|
|
|
def expl_from_problem(r, answers: List[str]) -> str:
|
|
"""Per-row explanation from the problem itself (no shared canned sentence)."""
|
|
tt = str(r.get("task_type", "")).strip() or "linguistics"
|
|
lang = str(r.get("task_lang", "")).strip() or "the target language"
|
|
q = re.sub(r"\s+", " ", str(r.get("query", "")).strip())
|
|
q = q[:90] + ("…" if len(q) > 90 else "")
|
|
n = len(answers)
|
|
preview = ", ".join(a for a in answers[:3] if a)
|
|
if len(answers) > 3:
|
|
preview += ", …"
|
|
bit = f" yielding {preview}" if preview else ""
|
|
return (
|
|
f"Solved this {tt} item set ({n} answers) in {lang} from the given "
|
|
f"examples, then applied the pattern to: {q}{bit}."
|
|
)
|
|
|
|
|
|
def clean_expl(text: str, fallback: str) -> str:
|
|
text = re.sub(r"\s+", " ", (text or "").strip())
|
|
text = re.sub(r"(?i)^(explanation|reasoning)\s*[:=\-]\s*", "", text).strip()
|
|
if not text:
|
|
return fallback
|
|
m = re.match(r"(.+?[.!?])(?:\s|$)", text)
|
|
if m:
|
|
text = m.group(1).strip()
|
|
if len(text) > 320:
|
|
text = text[:320].rsplit(" ", 1)[0].strip() + "."
|
|
return text or fallback
|
|
|
|
|
|
def save(rows: list[dict]) -> None:
|
|
import pandas as pd
|
|
|
|
cols = ["id", "pred"] + (["explanation"] if WRITE_EXPLANATIONS else [])
|
|
pd.DataFrame(rows, columns=cols).to_csv(OUTPUT_CSV, index=False)
|
|
|
|
|
|
def encode(tok, messages, device):
|
|
ids = tok.apply_chat_template(
|
|
messages, add_generation_prompt=True, return_tensors="pt"
|
|
)
|
|
if hasattr(ids, "input_ids"):
|
|
ids = ids["input_ids"]
|
|
return ids.to(device)
|
|
|
|
|
|
def gen(tok, model, ids, max_new: int) -> str:
|
|
import torch
|
|
|
|
with torch.no_grad():
|
|
out = model.generate(
|
|
ids,
|
|
max_new_tokens=max_new,
|
|
do_sample=False,
|
|
pad_token_id=tok.eos_token_id,
|
|
)
|
|
return tok.decode(out[0][ids.shape[-1] :], skip_special_tokens=True).strip()
|
|
|
|
|
|
def main() -> None:
|
|
import pandas as pd
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
started = time.monotonic()
|
|
tok = AutoTokenizer.from_pretrained(MODEL_ID)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID, torch_dtype=torch.float16, device_map="auto"
|
|
).eval()
|
|
device = next(model.parameters()).device
|
|
print(f"v8+expl loaded in {time.monotonic()-started:.1f}s", flush=True)
|
|
|
|
df = pd.read_csv(INPUT_CSV, dtype=str).fillna("")
|
|
n = len(df)
|
|
rows: list[dict] = []
|
|
meta: list[dict] = [] # context for phase-2 expl
|
|
durations: list[float] = []
|
|
|
|
# ----- phase 1: exact v8 answers -----
|
|
for _, r in df.iterrows():
|
|
elapsed = time.monotonic() - started
|
|
remaining = n - len(rows)
|
|
if elapsed >= SOFT_DEADLINE or (
|
|
remaining > 1 and elapsed + remaining * 8 > SOFT_DEADLINE + 30
|
|
):
|
|
answers: List[str] = []
|
|
row = {"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)}
|
|
if WRITE_EXPLANATIONS:
|
|
row["explanation"] = expl_from_problem(r, answers)
|
|
rows.append(row)
|
|
meta.append({"r": r, "answers": answers})
|
|
save(rows)
|
|
print(f"{len(rows)}/{n} DEADLINE", flush=True)
|
|
continue
|
|
|
|
task_type = str(r.get("task_type", "")).strip()
|
|
messages = [
|
|
{"role": "system", "content": SYSTEM_V8},
|
|
{
|
|
"role": "user",
|
|
"content": f"{r['context'].strip()}\n\n{r['query'].strip()}",
|
|
},
|
|
]
|
|
ids = encode(tok, messages, device)
|
|
prompt_len = ids.shape[-1]
|
|
|
|
t0 = time.monotonic()
|
|
text = gen(tok, model, ids, MAX_NEW)
|
|
# v8 parse: all non-empty lines (identical to submission8)
|
|
answers = [ln.strip() for ln in text.splitlines() if ln.strip()]
|
|
answers = safe_normalize_answers(answers, task_type)
|
|
durations.append(time.monotonic() - t0)
|
|
|
|
row = {"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)}
|
|
if WRITE_EXPLANATIONS:
|
|
row["explanation"] = expl_from_problem(r, answers)
|
|
rows.append(row)
|
|
meta.append({"r": r, "answers": answers})
|
|
save(rows)
|
|
print(
|
|
f"{len(rows)}/{n} ans {durations[-1]:.1f}s task={task_type or '?'}",
|
|
flush=True,
|
|
)
|
|
|
|
# ----- phase 2: model explanations with leftover time only -----
|
|
expl_model = 0
|
|
if WRITE_EXPLANATIONS:
|
|
for i, m in enumerate(meta):
|
|
time_left = SOFT_DEADLINE - (time.monotonic() - started)
|
|
if time_left < EXPL_STOP_LEFT:
|
|
break
|
|
r = m["r"]
|
|
answers = m["answers"]
|
|
if not answers:
|
|
continue
|
|
fb = expl_from_problem(r, answers)
|
|
preview = "; ".join(answers[:6])
|
|
if len(answers) > 6:
|
|
preview += "; ..."
|
|
expl_msgs = [
|
|
{
|
|
"role": "system",
|
|
"content": (
|
|
"Write ONE short English sentence on the main linguistic "
|
|
"rule used. No answers list, no preamble."
|
|
),
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": (
|
|
f"Task: {str(r.get('task_type', '')).strip()}\n"
|
|
f"Lang: {str(r.get('task_lang', '')).strip()}\n"
|
|
f"Answers: {preview}\n"
|
|
f"Problem:\n{str(r.get('context', ''))[:700]}\n\n"
|
|
f"{str(r.get('query', ''))[:400]}"
|
|
),
|
|
},
|
|
]
|
|
rows[i]["explanation"] = clean_expl(
|
|
gen(tok, model, encode(tok, expl_msgs, device), EXPL_MAX_NEW),
|
|
fb,
|
|
)
|
|
expl_model += 1
|
|
if expl_model % 10 == 0 or expl_model == 1:
|
|
save(rows)
|
|
print(
|
|
f"expl {expl_model}/{n} left={time_left:.0f}s",
|
|
flush=True,
|
|
)
|
|
|
|
save(rows)
|
|
print(
|
|
f"wrote {OUTPUT_CSV} | v8+expl | model_expl={expl_model}/{n} | "
|
|
f"total={time.monotonic()-started:.1f}s",
|
|
flush=True,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|