#import subprocess, sys
#subprocess.run([sys.executable, "-m", "pip", "install", "-q",
#                "transformers>=4.43", "accelerate>=0.30", "torch>=2.2", "pandas"], check=True)

import os
# The repo is the working directory at run time, and the sandbox has no network.
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
MODEL_ID = "."   # "." "./model"   # your model's weights are shipped in this repo

import re
import json
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
tok = AutoTokenizer.from_pretrained(MODEL_ID)

if tok.pad_token_id is None:
    tok.pad_token = tok.eos_token

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 = []


# ===== CHANGE HERE — 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."
)

# ===== CHANGE HERE — 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


for _, r in df.iterrows():
    messages = [
        {
            "role": "system",
            "content": SYSTEM
        },
        {
            "role": "user",
            "content": f"{r['context'].strip()}\n\n{r['query'].strip()}",
        },
    ]

    inputs = tok.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    ).to(model.device)

    input_length = inputs["input_ids"].shape[-1]

    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=512,
            do_sample=False,
            pad_token_id=tok.pad_token_id,
        )

    text = tok.decode(
        out[0][input_length:],
        skip_special_tokens=True,
    ).strip()

    answers = parse_answers(text)
    
    #answers = [
    #    line.strip()
    #    for line in text.splitlines()
    #    if line.strip()
    #]

    rows.append({
        "id": r["id"],
        "pred": json.dumps(answers, ensure_ascii=False),
    })

pd.DataFrame(rows).to_csv("submission.csv", index=False)