import os, subprocess, sys, time
T0 = time.time()
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
MODEL_ID = "."
DEADLINE = 26 * 60          # write CSV by 26 min, before the 30-min kill
MAX_NEW = 320               # answers are short; don't burn clock on 512

subprocess.run([sys.executable, "-m", "pip", "install", "-q", "bitsandbytes"], check=True)
print(f"[setup] deps ok at {time.time()-T0:.0f}s", flush=True)

import json
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                         bnb_4bit_compute_dtype=torch.float16)
tok = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, quantization_config=bnb, device_map="auto"
).eval()
print(f"[model] loaded at {time.time()-T0:.0f}s", flush=True)

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

rows = []
for i, (_, r) in enumerate(df.iterrows()):
    if time.time() - T0 > DEADLINE:
        print(f"[deadline] stopping at {i}/{len(df)}", flush=True)
        rows.extend({"id": rr["id"], "pred": json.dumps([])}
                    for _, rr in df.iloc[i:].iterrows())
        break
    messages = [
        {"role": "system", "content":
            "You solve International Linguistics Olympiad problems. Answer every numbered "
            "item. Put each answer on its own line, in order, with no numbering, no labels, "
            "and no extra text. Answer in the language the item asks for."},
        {"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, do_sample=False,
                             pad_token_id=tok.eos_token_id)
    text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
    answers = [ln.strip() for ln in text.splitlines() if ln.strip()]
    rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
    print(f"{i+1}/{len(df)} done at {time.time()-T0:.0f}s", flush=True)

pd.DataFrame(rows).to_csv("submission.csv", index=False)
print(f"wrote submission.csv at {time.time()-T0:.0f}s", flush=True)