import os
# The repo is the working directory at run time, and there is no network.
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
MODEL_ID = "."
MAX_NEW_TOKENS = 1536   # room to reason; lower = faster but answers may get cut off

import re
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()

# ===== your prompt (the main lever: how you ask the model) =====
SYSTEM = (
    "You solve International Linguistics Olympiad problems by reasoning from the "
    "data you are given. You may meet a task type you have never seen: read the "
    "instruction and the examples, and answer in the same form they use. "
    "Common task types and what to give -- "
    "translation: the translated form only, in the language the task asks for; "
    "fill_blanks: only the missing form for each blank; "
    "match_letters: only the option letter (for example A, B, C); "
    "text_to_num: the number in digits; "
    "num_to_text: the number written out in words, in the language asked; "
    "any other type: give exactly what the instruction asks, nothing else. "
    "Reason step by step first. Then write a line that says exactly FINAL ANSWERS: "
    "and, below it, one answer per line in the order the items are asked -- the "
    "bare answer only, no numbering, no quotes, no extra text."
)

# ===== how you read the answers back (must match the format your prompt asks for) =====
def parse_answers(text):
    """Keep only the lines after the last 'FINAL ANSWERS:' marker, one answer per line."""
    marker = list(re.finditer(r"(?im)^\s*final answers?\s*:?\s*$", text))
    if marker:
        text = text[marker[-1].end():]
    answers = []
    for line in text.splitlines():
        line = re.sub(r"^\s*\d+[.)]\s*", "", line).strip()   # drop "1. " / "2) " if the model adds it
        if line:
            answers.append(line)
    return answers

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

rows = []
for _, r in df.iterrows():
    messages = [
        {"role": "system", "content": SYSTEM},
        {"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=MAX_NEW_TOKENS, do_sample=False)
    text = tok.decode(out[0][enc["input_ids"].shape[-1]:], skip_special_tokens=True).strip()
    answers = parse_answers(text)
    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)