deepseek_r1_distill_qwen_7b_awq_rag/script.py

#!/usr/bin/env python3
"""Offline DeepSeek-R1-Distill-Qwen-7B-AWQ book-RAG IOL-AI submission."""

from __future__ import annotations

import argparse
import csv
import io
import json
import os
import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any


# The evaluation container has no internet access. Fail locally instead of waiting
# for network retries if a model/tokenizer file was not included in the repository.
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
os.environ.setdefault("HF_DATASETS_OFFLINE", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

REPO_DIR = Path(__file__).resolve().parent
RESOURCE_DIR = REPO_DIR / "rag_resources"
DEFAULT_INPUT = Path(os.environ.get("IOL_TEST_CSV", "/tmp/data/test.csv"))
DEFAULT_OUTPUT = Path(os.environ.get("IOL_SUBMISSION_CSV", "submission.csv"))

from rag_resources.retriever import BookRetriever  # noqa: E402


@dataclass
class ParseResult:
    thinking_trace: str
    final_text: str
    answers: list[str]
    valid: bool
    error: str = ""


def _load_config() -> dict[str, Any]:
    with (RESOURCE_DIR / "config.json").open("r", encoding="utf-8") as handle:
        config = json.load(handle)

    integer_overrides = {
        "IOL_TOP_METHODS": "top_methods",
        "IOL_TOP_EXAMPLES": "top_examples",
        "IOL_RAG_MAX_CHARS": "rag_max_chars",
        "IOL_MAX_REASONING_TOKENS": "max_reasoning_tokens",
        "IOL_MAX_ANSWER_TOKENS": "max_answer_tokens",
        "IOL_ANSWER_RETRY_TOKENS": "answer_retry_tokens",
        "IOL_EXPLANATION_MAX_NEW_TOKENS": "explanation_max_new_tokens",
        "IOL_SEED": "seed",
    }
    for environment_name, config_name in integer_overrides.items():
        if environment_name in os.environ:
            config[config_name] = int(os.environ[environment_name])
    if "IOL_CHAR_TFIDF_WEIGHT" in os.environ:
        config["char_tfidf_weight"] = float(os.environ["IOL_CHAR_TFIDF_WEIGHT"])
    if "IOL_ENABLE_EXPLANATIONS" in os.environ:
        config["enable_explanations"] = os.environ[
            "IOL_ENABLE_EXPLANATIONS"
        ].strip().casefold() in {"1", "true", "yes", "on"}
    return config


def _clean_answer_line(line: str) -> str:
    value = line.strip()
    value = re.sub(r"^(?:[-*•]\s+)", "", value)
    value = re.sub(r"^(?:\(?\d+\)?|\(?[A-Za-z]\)?)[.):]\s+", "", value)
    value = value.strip()
    pairs = {'"': '"', "'": "'", "`": "`", "“": "”", "‘": "’"}
    if len(value) >= 2 and value[0] in pairs and value[-1] == pairs[value[0]]:
        value = value[1:-1].strip()
    return value


def parse_model_response(thinking_trace: str, final_text: str) -> ParseResult:
    """Parse a separate reasoning trace and a strict FINAL ANSWERS block."""
    marker_matches = list(
        re.finditer(r"(?im)^\s*FINAL\s+ANSWERS\s*:\s*", final_text)
    )
    if not marker_matches:
        return ParseResult(
            thinking_trace=thinking_trace.strip(),
            final_text=final_text.strip(),
            answers=[],
            valid=False,
            error="missing FINAL ANSWERS: marker",
        )

    marker = marker_matches[-1]
    reasoning_outside_think = final_text[:marker.start()].strip()
    combined_trace = "\n\n".join(
        part for part in (thinking_trace.strip(), reasoning_outside_think) if part
    )
    answer_block = final_text[marker.end():].strip()
    answer_block = re.sub(r"^```(?:text)?\s*", "", answer_block, flags=re.I)
    answer_block = re.sub(r"\s*```\s*$", "", answer_block)
    if not answer_block:
        return ParseResult(
            thinking_trace=combined_trace,
            final_text=final_text.strip(),
            answers=[],
            valid=False,
            error="empty FINAL ANSWERS block",
        )

    # Be tolerant if the model emits a JSON list even though bare lines were asked for.
    answers: list[str] = []
    parsed_json_list = False
    if answer_block.startswith("["):
        try:
            decoded = json.loads(answer_block)
            if isinstance(decoded, list):
                parsed_json_list = True
                answers = [str(value).strip() for value in decoded if str(value).strip()]
        except json.JSONDecodeError:
            pass

    if not answers and not parsed_json_list:
        for line in answer_block.splitlines():
            stripped = line.strip()
            if not stripped or stripped.startswith("```"):
                continue
            if re.match(r"(?i)^(?:explanation|reasoning|notes?)\s*:", stripped):
                break
            cleaned = _clean_answer_line(stripped)
            if cleaned:
                answers.append(cleaned)

    if not answers:
        return ParseResult(
            thinking_trace=combined_trace,
            final_text=final_text.strip(),
            answers=[],
            valid=False,
            error="no non-empty answers in FINAL ANSWERS block",
        )
    return ParseResult(
        thinking_trace=combined_trace,
        final_text=final_text.strip(),
        answers=answers,
        valid=True,
    )


def _split_deepseek_response(response_text: str) -> tuple[str, str]:
    """Separate a completed DeepSeek think block from its visible answer."""
    open_marker = "<think>"
    close_marker = "</think>"
    open_index = response_text.find(open_marker)
    close_index = response_text.rfind(close_marker)
    if open_index >= 0 and close_index > open_index:
        reasoning = response_text[open_index + len(open_marker):close_index].strip()
        visible = (
            response_text[:open_index] + response_text[close_index + len(close_marker):]
        ).strip()
        return reasoning, visible
    if open_index < 0 and close_index >= 0:
        # The solve prompt pre-fills <think>, so generated tokens commonly begin
        # with the reasoning content and contain only the closing marker.
        reasoning = response_text[:close_index].strip()
        visible = response_text[close_index + len(close_marker):].strip()
        return reasoning, visible
    # If generation ended before </think> but still emitted the required final
    # marker, leave the text intact so the strict parser can recover that block.
    return "", response_text.replace(open_marker, "", 1).strip()


def _version_tuple(version: str) -> tuple[int, int, int]:
    numbers = [int(value) for value in re.findall(r"\d+", version)[:3]]
    return tuple((numbers + [0, 0, 0])[:3])  # type: ignore[return-value]


def _keep_last_token_hidden_state(_module: Any, inputs: tuple[Any, ...]) -> tuple[Any, ...] | None:
    """Avoid materializing full-sequence vocabulary logits during generation."""
    if not inputs:
        return None
    hidden_states = inputs[0]
    if hidden_states.ndim == 3 and hidden_states.shape[1] > 1:
        return (hidden_states[:, -1:, :],) + inputs[1:]
    return None


def load_model_and_tokenizer() -> tuple[Any, Any, Any]:
    """Load the pre-quantized AWQ model shipped in this repository."""
    try:
        import torch
        import transformers
        from transformers import AutoModelForCausalLM, AutoTokenizer
    except ImportError as exc:
        raise RuntimeError(
            "PyTorch, Transformers, Accelerate, AutoAWQ, and safetensors must be "
            "available in the evaluation image."
        ) from exc

    if _version_tuple(transformers.__version__) < (4, 37, 0):
        raise RuntimeError(
            "DeepSeek-R1-Distill-Qwen-7B uses the Qwen2 architecture and requires "
            f"transformers>=4.37.0; found {transformers.__version__}."
        )
    model_config_path = REPO_DIR / "config.json"
    if not model_config_path.exists():
        raise RuntimeError(
            "DeepSeek-R1-Distill-Qwen-7B-AWQ files are missing. Put the complete model and "
            "tokenizer snapshot in the same repository directory as script.py."
        )
    with model_config_path.open("r", encoding="utf-8") as handle:
        model_metadata = json.load(handle)
    quantization_method = str(
        model_metadata.get("quantization_config", {}).get("quant_method", "")
    ).casefold()
    if (
        model_metadata.get("model_type") != "qwen2"
        or quantization_method != "awq"
        or int(model_metadata.get("hidden_size", 0)) != 3584
        or int(model_metadata.get("num_hidden_layers", 0)) != 28
    ):
        raise RuntimeError(
            "This script expects a 4-bit AWQ conversion of "
            "DeepSeek-R1-Distill-Qwen-7B (Qwen2, hidden_size=3584, 28 layers)."
        )

    tokenizer = AutoTokenizer.from_pretrained(
        str(REPO_DIR), local_files_only=True, trust_remote_code=False
    )
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id
    model = AutoModelForCausalLM.from_pretrained(
        str(REPO_DIR),
        local_files_only=True,
        trust_remote_code=False,
        device_map="auto",
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True,
    ).eval()
    model.config.use_cache = True
    # Transformers 4.44.1 projects every prompt token through Qwen2's large
    # vocabulary head and then upcasts all logits to FP32. Generation only uses
    # the final-token logits, so trim the token dimension immediately before
    # lm_head to avoid a multi-gigabyte prefill allocation on the T4.
    model.lm_head.register_forward_pre_hook(_keep_last_token_hidden_state)
    input_device = model.get_input_embeddings().weight.device
    return model, tokenizer, input_device


def _problem_prompt(row: dict[str, Any], rag_text: str, retry_note: str = "") -> str:
    work_language = str(row.get("work_lang", "")).strip()
    task_language = str(row.get("task_lang", "")).strip()
    language_lines = []
    if work_language:
        language_lines.append(f"Working language: {work_language}")
    if task_language:
        language_lines.append(f"Problem language: {task_language}")
    language_metadata = "\n".join(language_lines)

    sections = [rag_text] if rag_text else []
    sections.append("CURRENT PROBLEM")
    if language_metadata:
        sections.append(language_metadata)
    sections.extend(
        [
            "CONTEXT:\n" + str(row.get("context", "")),
            "QUERY:\n" + str(row.get("query", "")),
        ]
    )
    if retry_note:
        sections.append(retry_note)
    return "\n\n".join(sections)


def _encode_fitted_prompt(
    tokenizer: Any,
    system_prompt: str,
    row: dict[str, Any],
    rag_text: str,
    retry_note: str,
    context_window: int,
    requested_new_tokens: int,
) -> tuple[dict[str, Any], int]:
    """Trim retrieved material, never the current problem, to fit the context."""
    fitted_rag = rag_text
    while True:
        # DeepSeek recommends placing all instructions in the user message for
        # this R1 distillation rather than using a separate system message.
        messages = [
            {
                "role": "user",
                "content": system_prompt
                + "\n\n"
                + _problem_prompt(row, fitted_rag, retry_note),
            }
        ]
        rendered = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
        rendered += "<think>\n"
        encoded = tokenizer(rendered, return_tensors="pt", add_special_tokens=False)
        prompt_tokens = int(encoded["input_ids"].shape[-1])
        available = context_window - prompt_tokens
        if available >= requested_new_tokens:
            return encoded, requested_new_tokens
        if fitted_rag:
            if len(fitted_rag) <= 600:
                fitted_rag = ""
            else:
                new_length = max(0, int(len(fitted_rag) * 0.72))
                fitted_rag = fitted_rag[:new_length].rsplit("\n", 1)[0].rstrip()
                fitted_rag += "\n[retrieved material shortened to fit the model context]"
            continue
        if available >= 256:
            return encoded, available
        raise RuntimeError(
            f"The current problem and solver instructions use {prompt_tokens} tokens, "
            f"leaving only {available} tokens in the {context_window}-token context."
        )


def generate_once(
    model: Any,
    tokenizer: Any,
    input_device: Any,
    system_prompt: str,
    row: dict[str, Any],
    rag_text: str,
    retry_note: str,
    config: dict[str, Any],
    seed: int,
) -> ParseResult:
    import torch

    configured_reasoning_tokens = int(config["max_reasoning_tokens"])
    configured_answer_tokens = int(config["max_answer_tokens"])
    encoded, available_generation_tokens = _encode_fitted_prompt(
        tokenizer=tokenizer,
        system_prompt=system_prompt,
        row=row,
        rag_text=rag_text,
        retry_note=retry_note,
        context_window=int(config["model_context_tokens"]),
        requested_new_tokens=(
            configured_reasoning_tokens + configured_answer_tokens + 16
        ),
    )
    encoded = {name: value.to(input_device) for name, value in encoded.items()}
    answer_tokens = min(configured_answer_tokens, available_generation_tokens - 272)
    reasoning_tokens = min(
        configured_reasoning_tokens,
        available_generation_tokens - answer_tokens - 16,
    )
    if reasoning_tokens < 256 or answer_tokens < 256:
        raise RuntimeError(
            "The fitted prompt does not leave at least 256 tokens for both the "
            "reasoning and answer stages."
        )

    close_ids = tokenizer(
        "</think>", add_special_tokens=False, return_tensors="pt"
    )["input_ids"][0].tolist()
    if len(close_ids) != 1:
        raise RuntimeError("Expected </think> to be one tokenizer token.")
    close_id = int(close_ids[0])

    def sample(
        input_ids: Any,
        attention_mask: Any,
        max_tokens: int,
        sample_seed: int,
        eos_token_id: Any,
    ) -> Any:
        torch.manual_seed(sample_seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(sample_seed)
        with torch.inference_mode():
            return model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=max_tokens,
                do_sample=True,
                temperature=float(config["temperature"]),
                top_p=float(config["top_p"]),
                top_k=int(config["top_k"]),
                repetition_penalty=float(config["repetition_penalty"]),
                use_cache=True,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=eos_token_id,
            )

    reasoning_output = sample(
        encoded["input_ids"],
        encoded["attention_mask"],
        reasoning_tokens,
        seed,
        [tokenizer.eos_token_id, close_id],
    )
    prompt_length = int(encoded["input_ids"].shape[-1])
    reasoning_ids = reasoning_output[0, prompt_length:].tolist()
    reasoning_closed = bool(reasoning_ids and reasoning_ids[-1] == close_id)
    if reasoning_ids and reasoning_ids[-1] == tokenizer.eos_token_id:
        reasoning_ids = reasoning_ids[:-1]
        reasoning_output = reasoning_output[:, :-1]
    reasoning_text = tokenizer.decode(
        reasoning_ids, skip_special_tokens=True
    ).strip()
    thinking_trace, _ = _split_deepseek_response(reasoning_text)
    if not thinking_trace:
        thinking_trace = reasoning_text.replace("</think>", "").strip()
    print(
        f"  reasoning tokens={len(reasoning_ids)}/{reasoning_tokens}; "
        f"closed={'yes' if reasoning_closed else 'forced'}",
        flush=True,
    )

    answer_prefix = "\nFINAL ANSWERS:\n" if reasoning_closed else "</think>\nFINAL ANSWERS:\n"
    prefix_ids = tokenizer(
        answer_prefix, add_special_tokens=False, return_tensors="pt"
    )["input_ids"].to(input_device)
    answer_input_ids = torch.cat((reasoning_output, prefix_ids), dim=1)
    answer_attention_mask = torch.ones_like(answer_input_ids)
    answer_budgets = [answer_tokens, int(config["answer_retry_tokens"])]
    last_result: ParseResult | None = None
    for answer_attempt, answer_budget in enumerate(answer_budgets, start=1):
        answer_output = sample(
            answer_input_ids,
            answer_attention_mask,
            answer_budget,
            seed + answer_attempt,
            tokenizer.eos_token_id,
        )
        answer_ids = answer_output[0, answer_input_ids.shape[-1]:].tolist()
        answer_text = tokenizer.decode(
            answer_ids, skip_special_tokens=True
        ).strip()
        final_text = "FINAL ANSWERS:\n" + answer_text
        result = parse_model_response(thinking_trace, final_text)
        print(
            f"  answer attempt {answer_attempt}/{len(answer_budgets)}: "
            f"tokens={len(answer_ids)}/{answer_budget}; "
            f"parser={'accepted' if result.valid else result.error}",
            flush=True,
        )
        if result.valid:
            return result
        last_result = result

    assert last_result is not None
    return last_result


def solve_row(
    row: dict[str, Any],
    row_index: int,
    retriever: BookRetriever,
    model: Any,
    tokenizer: Any,
    input_device: Any,
    system_prompt: str,
    config: dict[str, Any],
) -> ParseResult:
    retrieval = retriever.retrieve(
        row,
        top_methods=int(config["top_methods"]),
        top_examples=int(config["top_examples"]),
        char_tfidf_weight=float(config["char_tfidf_weight"]),
    )
    rag_text = retriever.format_for_prompt(
        retrieval, max_chars=int(config["rag_max_chars"])
    )
    return generate_once(
        model=model,
        tokenizer=tokenizer,
        input_device=input_device,
        system_prompt=system_prompt,
        row=row,
        rag_text=rag_text,
        retry_note="",
        config=config,
        seed=int(config["seed"]) + row_index * 3,
    )


def generate_explanation(
    solution: ParseResult,
    model: Any,
    tokenizer: Any,
    input_device: Any,
    config: dict[str, Any],
) -> str:
    """Summarize the model's accepted reasoning for the optional jury track."""
    import torch

    source_reasoning = solution.thinking_trace.strip() or solution.final_text.strip()
    answer_text = "\n".join(solution.answers)
    messages = [
        {
            "role": "user",
            "content": (
                "Write a short, human-readable explanation of the solution in a few "
                "concise bullet points. State the linguistic rule or pattern found, "
                "the key evidence, and how the final answers follow. Do not reproduce "
                "the raw reasoning trace, trial and error, or meta-commentary. Do not "
                "change the final answers. Output only the explanation.\n\n"
                f"MODEL REASONING:\n{source_reasoning}\n\n"
                f"FINAL ANSWERS:\n{answer_text}"
            ),
        },
    ]
    rendered = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    encoded = tokenizer(rendered, return_tensors="pt", add_special_tokens=False)
    encoded = {name: value.to(input_device) for name, value in encoded.items()}
    prompt_tokens = int(encoded["input_ids"].shape[-1])
    available = int(config["model_context_tokens"]) - prompt_tokens
    max_new_tokens = min(int(config["explanation_max_new_tokens"]), available)
    if max_new_tokens < 32:
        return (
            "- Inferred the relevant linguistic patterns from the supplied examples.\n"
            "- Applied those patterns to produce the listed answers."
        )

    with torch.inference_mode():
        output = model.generate(
            **encoded,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            use_cache=True,
            pad_token_id=tokenizer.pad_token_id,
        )
    generated_ids = output[0, prompt_tokens:].tolist()
    raw_explanation = tokenizer.decode(
        generated_ids, skip_special_tokens=True
    ).strip()
    if "<think>" in raw_explanation and "</think>" not in raw_explanation:
        raw_explanation = ""
    _, explanation = _split_deepseek_response(raw_explanation)
    if explanation:
        return explanation
    return (
        "- Inferred the relevant linguistic patterns from the supplied examples.\n"
        "- Applied those patterns to produce the listed answers."
    )


def read_test_rows(path: Path) -> list[dict[str, Any]]:
    with path.open("r", encoding="utf-8-sig", newline="") as handle:
        return list(csv.DictReader(handle))


def _write_submission_stream(handle: Any, rows: list[dict[str, str]]) -> None:
    writer = csv.DictWriter(handle, fieldnames=["id", "pred", "explanation"])
    writer.writeheader()
    writer.writerows(rows)


def write_submission(path: Path, rows: list[dict[str, str]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary_path = path.with_name(f".{path.name}.tmp")
    with temporary_path.open("w", encoding="utf-8", newline="") as handle:
        _write_submission_stream(handle, rows)
    os.replace(temporary_path, path)


def run_self_test() -> None:
    import numpy as np

    hidden_states = np.arange(24).reshape(1, 3, 8)
    trimmed = _keep_last_token_hidden_state(None, (hidden_states,))
    assert trimmed is not None and trimmed[0].shape == (1, 1, 8)
    assert np.array_equal(trimmed[0], hidden_states[:, -1:, :])
    assert _keep_last_token_hidden_state(None, (hidden_states[:, -1:, :],)) is None

    parsed = parse_model_response(
        "A useful analysis.",
        "One last check.\nFINAL ANSWERS:\n1. čha\n2) multi word form",
    )
    assert parsed.valid
    assert parsed.answers == ["čha", "multi word form"]
    assert "One last check." in parsed.thinking_trace
    assert not parse_model_response("", "FINAL ANSWERS:\n").valid
    assert not parse_model_response("", "FINAL ANSWERS:\n[]").valid
    assert not parse_model_response("", 'FINAL ANSWERS:\n[""]').valid
    assert not parse_model_response("", "The answer is x.").valid

    deepseek_trace, deepseek_final = _split_deepseek_response(
        "<think>Compare the recurring suffixes.</think>\nFINAL ANSWERS:\nform"
    )
    assert deepseek_trace == "Compare the recurring suffixes."
    assert deepseek_final == "FINAL ANSWERS:\nform"
    deepseek_parsed = parse_model_response(deepseek_trace, deepseek_final)
    assert deepseek_parsed.valid and deepseek_parsed.answers == ["form"]
    prefixed_trace, prefixed_final = _split_deepseek_response(
        "Compare the recurring suffixes.</think>\nFINAL ANSWERS:\nform"
    )
    assert prefixed_trace == "Compare the recurring suffixes."
    assert prefixed_final == "FINAL ANSWERS:\nform"

    retriever = BookRetriever(RESOURCE_DIR)
    canary_row = {
        "context": "The following words are number expressions in an unknown language.",
        "query": "Determine the rule and write the number 25 in words.",
        "answer": "PRIVATE_ANSWER_CANARY",
    }
    result = retriever.retrieve(
        canary_row, top_methods=2, top_examples=2, char_tfidf_weight=3.0
    )
    prompt = retriever.format_for_prompt(result, max_chars=8000)
    assert result["methods"] and result["examples"]
    assert "PRIVATE_ANSWER_CANARY" not in prompt
    assert all(example.get("string_only_usable") for example in result["examples"])

    exact_example = retriever.examples[0]
    exact_result = retriever.retrieve(
        {"context": exact_example["context"], "query": exact_example["query"]},
        top_methods=1,
        top_examples=1,
        char_tfidf_weight=3.0,
    )
    assert exact_result["examples"][0]["id"] == exact_example["id"]

    submission_buffer = io.StringIO(newline="")
    _write_submission_stream(
        submission_buffer,
        [
            {
                "id": "007",
                "pred": json.dumps(["čha", "multi word"], ensure_ascii=False),
                "explanation": "- Identified the relevant pattern.",
            }
        ],
    )
    submission_buffer.seek(0)
    assert list(csv.DictReader(submission_buffer)) == [
        {
            "id": "007",
            "pred": '["čha", "multi word"]',
            "explanation": "- Identified the relevant pattern.",
        }
    ]

    incremental_rows = [
        {"id": "001", "pred": "[]", "explanation": ""},
        {"id": "002", "pred": "[]", "explanation": ""},
    ]
    incremental_rows[0]["pred"] = json.dumps(["answer"], ensure_ascii=False)
    incremental_buffer = io.StringIO(newline="")
    _write_submission_stream(incremental_buffer, incremental_rows)
    incremental_buffer.seek(0)
    assert list(csv.DictReader(incremental_buffer)) == incremental_rows
    print(
        "Self-test passed: DeepSeek think parsing, last-token logits hook, strict "
        "answer parsing, hybrid book-only retrieval, exact-match retrieval, and "
        "submission CSV serialization."
    )


def main() -> None:
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8")
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, default=DEFAULT_INPUT)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument(
        "--explanations",
        choices=("on", "off"),
        default=None,
        help="Generate optional jury-track explanations (default: config setting)",
    )
    parser.add_argument(
        "--self-test", action="store_true", help="Test parser/retrieval without loading the model"
    )
    args = parser.parse_args()

    if args.self_test:
        run_self_test()
        return
    if not args.input.exists():
        raise SystemExit(f"Test file not found: {args.input}")

    config = _load_config()
    explanations_enabled = (
        args.explanations == "on"
        if args.explanations is not None
        else bool(config.get("enable_explanations", False))
    )
    system_prompt = (RESOURCE_DIR / "system_prompt.txt").read_text(encoding="utf-8").strip()
    test_rows = read_test_rows(args.input)
    retriever = BookRetriever(RESOURCE_DIR)
    model, tokenizer, input_device = load_model_and_tokenizer()
    submission_rows: list[dict[str, str]] = [
        {
            "id": str(row.get("id", "")),
            "pred": "[]",
            "explanation": "",
        }
        for row in test_rows
    ]
    write_submission(args.output, submission_rows)
    print(
        f"Initialized incremental submission with {len(submission_rows)} rows at "
        f"{args.output}; explanations={'on' if explanations_enabled else 'off'}",
        flush=True,
    )

    for index, row in enumerate(test_rows):
        row_id = str(row.get("id", ""))
        print(f"[{index + 1}/{len(test_rows)}] solving id={row_id}", flush=True)
        solution = solve_row(
            row=row,
            row_index=index,
            retriever=retriever,
            model=model,
            tokenizer=tokenizer,
            input_device=input_device,
            system_prompt=system_prompt,
            config=config,
        )
        explanation = ""
        if explanations_enabled:
            explanation = generate_explanation(
                solution=solution,
                model=model,
                tokenizer=tokenizer,
                input_device=input_device,
                config=config,
            )
        submission_rows[index] = {
            "id": row_id,
            "pred": json.dumps(solution.answers, ensure_ascii=False),
            "explanation": explanation,
        }
        write_submission(args.output, submission_rows)
        print(f"Checkpointed prediction {index + 1}/{len(test_rows)}", flush=True)

    print(f"Wrote {len(submission_rows)} predictions to {args.output}", flush=True)


if __name__ == "__main__":
    main()