import os
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
MODEL_ID = "."
MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"

import time, json, re
import pandas as pd, torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MAX_NEW_TOKENS = 1024
t0 = time.time()
tok = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float16, device_map="auto").eval()
print(f"[load] running {MODEL_NAME} from repo weights ({MODEL_ID})", flush=True)
print(f"[load] model ready in {time.time()-t0:.0f}s", flush=True)

df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")

def extract_answers(text):
    # Prefer the section after "FINAL ANSWERS:" when present
    m = list(re.finditer(r'(?im)^[\s>*#-]*final answers?\s*[:.]?\s*$', text))
    if m:
        text = text[m[-1].end():]

    # 1) answers on lines like [answer] (model asked to put only final answers there)
    br = []
    for ln in text.splitlines():
        ln = ln.strip()
        m = re.match(r'^\[(.+)\]$', ln)
        if m:
            br.append(m.group(1).strip())
    if br:
        return br

    # 2) fallback: one cleaned answer per line
    out = []
    for ln in text.splitlines():
        ln = re.sub(r'^[\s>*#-]+', '', ln)
        ln = re.sub(r'^\d+[.)]\s*', '', ln).strip().strip("[]").strip()
        if ln:
            out.append(ln)
    return out

SYSTEM = ("You solve International Linguistics Olympiad problems by reasoning from the data given. "
    "You may face a task type you have never seen — read the instruction and adapt. Answer in the "
    "language the task asks for; for matching items give the option letter, for number items give "
    "digits or the written-out number as asked. First reason briefly. Then write your FINAL ANSWERS: "
    "one per item, in the order the items appear, each wrapped in square brackets like "
    "[answer], and nothing else on those lines.")

rows = []
for i, r in df.iterrows():
    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():
        out = model.generate(ids, max_new_tokens=MAX_NEW_TOKENS, do_sample=False)
    text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
    answers = extract_answers(text)
    rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
    print(f"[{i+1}/{len(df)}] id={r['id']} -> {len(answers)} answers", flush=True)

pd.DataFrame(rows).to_csv("submission.csv", index=False)
print(f"[done] wrote submission.csv ({len(rows)} rows) in {time.time()-t0:.0f}s", flush=True)