ratiocine/script.py

"""
Competition submission script — copy this into your HF repo as script.py.

The eval sandbox:
  - mounts the test set at /tmp/data/test.csv
  - has no internet
  - runs on a T4 (16GB)
  - has 30 minutes
  - has bitsandbytes and autoawq pre-installed

Strategy: Ship Qwen2.5-14B-Instruct-AWQ with HYBRID prompting.
The 14B-AWQ is the proven competition baseline (0.123 score). We improve
on the baseline with:
1. Task-specific CoT prompts for translation/fill_blanks (improves EM)
2. Direct prompts for match_letters/text_to_num/num_to_text (faster)
3. Adaptive max_new_tokens per task type (512 for CoT, 256 for direct)
4. Fixed answer parser (v1 dropped ~5% of correct answers)
5. Explanation column for human jury track
6. Time guard to never exceed 30-min limit

Why hybrid: Pure CoT was too slow (70s/problem) and exceeded the time
budget. Pure direct prompting gave EM=0.025 on the hidden test set.
CoT for hard tasks (translation, fill_blanks) improves exact matches by
letting the model reason carefully; direct prompting is fine for
pattern-matching tasks where reasoning doesn't help.
"""

import json
import os
import re
import time

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

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

from prompts import (
    get_system_prompt,
    USER_TEMPLATE,
    parse_answers,
    extract_analysis,
    count_query_items,
)

# Time budget: 30 min total. Reserve 3 min for model loading + CSV write.
TIME_BUDGET_S = 27 * 60  # 27 minutes for inference


def load_model():
    """Load the AWQ-quantized model for T4 16GB."""
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        device_map="auto",
        trust_remote_code=True,
        torch_dtype=torch.float16,
    )
    model.eval()
    print("[submit] Loaded Qwen2.5-14B-Instruct-AWQ", flush=True)
    return tokenizer, model


def solve_problem(
    tokenizer,
    model,
    context: str,
    query: str,
    task_type: str = "",
    max_new_tokens: int = 256,
) -> tuple[list[str], str]:
    """Generate answers for one IOL problem.

    For translation/fill_blanks: CoT reasoning (max 512 tokens).
    For match_letters/text_to_num/num_to_text: direct (256 tokens).
    """
    n_items = count_query_items(query)
    system_prompt = get_system_prompt(task_type)

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": USER_TEMPLATE.format(
            context=context.strip(), query=query.strip()
        )},
    ]

    text = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, tokenize=False
    )
    inputs = tokenizer(text, return_tensors="pt")
    input_ids = inputs["input_ids"].to(model.device)

    with torch.no_grad():
        # Greedy decoding (do_sample=False) — reproducible and best for our use case.
        # Beam search (num_beams=2) tested but caused catastrophic failures (0.0 score).
        out = model.generate(
            input_ids,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )

    generated = tokenizer.decode(
        out[0][input_ids.shape[-1]:], skip_special_tokens=True
    ).strip()

    answers = parse_answers(generated, n_expected=n_items, task_type=task_type)
    explanation = extract_analysis(generated)

    return answers, explanation


def _format_pred(answers: list[str]) -> str:
    """Format predictions for submission.

    Output is JSON-encoded list of answer strings (the IOL competition
    evaluator parses this with ast.literal_eval). We also include a
    pipe-separated fallback in a comment-like column for safety.
    """
    return json.dumps(answers, ensure_ascii=False)


def main():
    t_start = time.time()

    print("[submit] Loading model...", flush=True)
    tokenizer, model = load_model()

    print("[submit] Reading test set...", flush=True)
    df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
    n_problems = len(df)
    print(f"[submit] Loaded {n_problems} problems", flush=True)

    # Estimate time per problem type for adaptive budget
    # CoT tasks: ~20s each (512 tokens); direct tasks: ~5s each (256 tokens)
    COT_TASKS = {"translation", "fill_blanks"}
    DIRECT_TASKS = {"match_letters", "text_to_num", "num_to_text"}
    SHORT_TASKS = {"match_letters", "text_to_num"}  # single char / digits

    rows = []
    for idx, row in df.iterrows():
        elapsed = time.time() - t_start
        remaining = TIME_BUDGET_S - elapsed
        problems_left = n_problems - idx
        task_type = row.get("task_type", "")

        # Adaptive max_new_tokens based on time remaining and task type
        # Target: average ~10s per problem to fit 160 problems in 27 min
        # Budget per problem: 27*60 / 160 = 10.1s
        # CoT is ~17s, direct is ~5-8s, short is ~4s
        # Be aggressive: switch to fallback when remaining < 8s/problem
        if remaining < problems_left * 8 and remaining > 0 and idx > 0:
            # Tight on time — minimal tokens, direct mode (no CoT)
            current_max = 96
            use_cot = False
            if idx % 10 == 0:
                print(f"[submit] FAST MODE at {idx+1}/{n_problems} "
                      f"({remaining:.0f}s left, {remaining/problems_left:.1f}s/problem)",
                      flush=True)
        elif remaining < problems_left * 12 and remaining > 0 and idx > 0 and task_type in COT_TASKS:
            # Getting tight on CoT problems — reduce CoT max tokens
            current_max = 256
            use_cot = True
            if idx % 10 == 0:
                print(f"[submit] COOL DOWN at {idx+1}/{n_problems} "
                      f"({remaining:.0f}s left, {remaining/problems_left:.1f}s/problem)",
                      flush=True)
        elif task_type in COT_TASKS:
            # Verbose CoT (matches 0.0872 baseline that scored highest).
            # 512 tokens gives room for full step-by-step reasoning.
            current_max = 512
            use_cot = True
        elif task_type in SHORT_TASKS:
            # Short answers (single letters or digits) — keep tight
            current_max = 128
            use_cot = False
        else:
            # Direct for easy tasks (num_to_text, etc.)
            current_max = 256
            use_cot = False

        try:
            # If we need to force direct mode for time, swap to default prompt
            if not use_cot and task_type in COT_TASKS and remaining < problems_left * 8:
                # Switch to default prompt (direct) for time-constrained CoT tasks
                from prompts import _DEFAULT_PROMPT
                original_prompt = get_system_prompt(task_type)
                # Use default prompt via monkey-patch
                import prompts
                prompts._PROMPTS[task_type] = _DEFAULT_PROMPT
                answers, explanation = solve_problem(
                    tokenizer, model,
                    context=row["context"],
                    query=row["query"],
                    task_type=task_type,
                    max_new_tokens=current_max,
                )
                prompts._PROMPTS[task_type] = original_prompt
            else:
                answers, explanation = solve_problem(
                    tokenizer, model,
                    context=row["context"],
                    query=row["query"],
                    task_type=task_type,
                    max_new_tokens=current_max,
                )
        except Exception as e:
            print(f"[submit] ERROR at {idx+1}/{n_problems}: {e}", flush=True)
            n_items = count_query_items(row.get("query", ""))
            answers = [""] * max(n_items, 1)
            explanation = ""

        rows.append({
            "id": row["id"],
            "pred": _format_pred(answers),
            "explanation": explanation,
        })

        if (idx + 1) % 10 == 0 or idx == 0:
            print(f"[submit] {idx + 1}/{n_problems} done "
                  f"({elapsed:.0f}s elapsed, task={task_type})", flush=True)

    output = pd.DataFrame(rows)
    # Write to the path the eval system expects
    import os as _os
    _os.makedirs("/tmp/model", exist_ok=True)
    output.to_csv("/tmp/model/submission.csv", index=False)
    # Also write a backup at the relative path (in case CWD is /tmp/model)
    output.to_csv("submission.csv", index=False)
    total_elapsed = time.time() - t_start
    print(f"[submit] wrote submission.csv ({len(rows)} problems, "
          f"{total_elapsed:.0f}s total, cwd={_os.getcwd()})", flush=True)


if __name__ == "__main__":
    main()