script.py

import os
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

import re
import json
import time
import sys
import subprocess

try:
    import bitsandbytes  # noqa: F401
except ImportError:
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", "bitsandbytes"], check=True)
    import bitsandbytes

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

# Use the local model directory (uploaded to HF repo root)
MODEL_ID = "."
TIME_LIMIT_S = 28 * 60
START = time.time()


def time_left():
    return TIME_LIMIT_S - (time.time() - START)


def load_model():
    """Load model with 4-bit quantization"""
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True,
    )
    
    tok = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        quantization_config=bnb_config,
        device_map="auto",
        torch_dtype=torch.float16,
        local_files_only=True,
    ).eval()
    return tok, model


def generate(messages, max_new_tokens=3500):
    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,
            temperature=None,
            top_p=None,
            pad_token_id=tok.eos_token_id,
        )
    return tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()


def count_items(query: str) -> int:
    nums = re.findall(r"(?m)^\s*(\d+)\s*[.)]\s", query)
    if nums:
        return max(int(n) for n in nums)
    lines = [l for l in query.splitlines() if l.strip()]
    return max(1, len(lines))


def extract_json_list(text: str):
    text = text.strip()
    try:
        val = json.loads(text)
        if isinstance(val, list):
            return [str(x) for x in val]
    except Exception:
        pass
    m = re.search(r"\[.*\]", text, re.DOTALL)
    if m:
        try:
            val = json.loads(m.group(0))
            if isinstance(val, list):
                return [str(x) for x in val]
        except Exception:
            pass
    return None


def extract_answers_fallback(text: str, n: int):
    lines = [l.strip() for l in text.splitlines() if l.strip()]
    cleaned = []
    for l in lines:
        l = re.sub(r"^\s*\d+[.)]\s*", "", l)
        l = l.strip(" -\t")
        if l:
            cleaned.append(l)
    if len(cleaned) >= n:
        return cleaned[:n]
    parts = [p.strip() for p in text.replace("\n", ",").split(",") if p.strip()]
    if len(parts) >= n:
        return parts[:n]
    cleaned = cleaned or parts or [text.strip()]
    while len(cleaned) < n:
        cleaned.append(cleaned[-1] if cleaned else "")
    return cleaned[:n]


def fit_answers(answers, n):
    answers = list(answers)
    if len(answers) < n:
        answers = answers + [answers[-1] if answers else ""] * (n - len(answers))
    return answers[:n]


REASONING_SYSTEM = (
    "You are an expert at International Linguistics Olympiad (IOL) problems. "
    "You will be given a self-contained puzzle about an unfamiliar language: some example "
    "data (context) and a query asking you to translate, match, fill in blanks, or convert "
    "numbers. Work out the underlying grammar/vocabulary rules carefully and systematically "
    "from the given examples only. Show your step-by-step reasoning: identify morphemes, "
    "patterns, correspondences, and test your hypothesis against every example before "
    "answering the query."
)

EXTRACT_SYSTEM = (
    "You convert a linguistics-puzzle solution into a strict output format. "
    "Given the original problem and a reasoning trace, output ONLY a single JSON object "
    "with exactly two keys:\n"
    '  "answers": a JSON list of strings, one per numbered item in the query, in order.\n'
    '  "explanation": a short (3-6 bullet points, plain text with \n between bullets) '
    "human-readable summary of the key rules used to derive the answers. This is NOT the "
    "full reasoning trace, just a concise justification a person can read in under a minute.\n"
    "Output nothing except that JSON object."
)


def solve_row(context: str, query: str, n_items: int):
    problem_text = f"{context.strip()}\n\n{query.strip()}"

    reasoning = ""
    if time_left() > 60:
        try:
            reasoning = generate(
                [
                    {"role": "system", "content": REASONING_SYSTEM},
                    {"role": "user", "content": problem_text},
                ],
                max_new_tokens=3500,
            )
        except Exception as e:
            reasoning = f"(reasoning pass failed: {e})"

    answers, explanation = None, ""
    if time_left() > 30:
        extract_prompt = (
            f"PROBLEM:\n{problem_text}\n\n"
            f"REASONING TRACE:\n{reasoning}\n\n"
            f"The query has {n_items} numbered item(s). Now output the JSON object as instructed."
        )
        try:
            raw = generate(
                [
                    {"role": "system", "content": EXTRACT_SYSTEM},
                    {"role": "user", "content": extract_prompt},
                ],
                max_new_tokens=800,
            )
            try:
                obj = json.loads(raw.strip())
            except Exception:
                m = re.search(r"\{.*\}", raw, re.DOTALL)
                obj = json.loads(m.group(0)) if m else None
            if obj and isinstance(obj, dict):
                a = obj.get("answers")
                if isinstance(a, list):
                    answers = [str(x) for x in a]
                explanation = str(obj.get("explanation", "")).strip()
        except Exception:
            pass

    if answers is None:
        candidate = extract_json_list(reasoning)
        answers = candidate if candidate is not None else extract_answers_fallback(reasoning, n_items)
    if not explanation:
        snippet = reasoning.strip().replace("\n", " ")
        explanation = (snippet[:400] + "...") if len(snippet) > 400 else snippet

    answers = fit_answers(answers, n_items)
    return answers, explanation


def main():
    # Load model once
    global tok, model
    tok, model = load_model()
    
    df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
    out_rows = []

    for i, r in df.iterrows():
        rid = r["id"]
        context, query = r.get("context", ""), r.get("query", "")
        n_items = count_items(query)

        if time_left() < 45:
            answers = [""] * n_items
            explanation = ""
        else:
            try:
                answers, explanation = solve_row(context, query, n_items)
            except Exception as e:
                answers = [""] * n_items
                explanation = f"(error: {e})"

        out_rows.append(
            {
                "id": rid,
                "pred": json.dumps(answers, ensure_ascii=False),
                "explanation": explanation,
            }
        )
        print(f"[{i + 1}/{len(df)}] id={rid} n_items={n_items} time_left={time_left():.0f}s", flush=True)

    pd.DataFrame(out_rows, columns=["id", "pred", "explanation"]).to_csv(
        "submission.csv", index=False
    )
    print("wrote submission.csv", flush=True)


if __name__ == "__main__":
    main()