665 lines
24 KiB
Python
665 lines
24 KiB
Python
#!/usr/bin/env python3
|
||
"""Offline Qwen2.5-14B-AWQ + book-RAG submission for the IOL-AI challenge."""
|
||
|
||
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_NEW_TOKENS": "max_new_tokens",
|
||
"IOL_EXPLANATION_MAX_NEW_TOKENS": "explanation_max_new_tokens",
|
||
"IOL_MAX_ATTEMPTS": "max_attempts",
|
||
"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 _best_effort_answers(final_texts: list[str]) -> list[str]:
|
||
"""Keep submission shape valid if all strict retries fail."""
|
||
for text in reversed(final_texts):
|
||
candidates = []
|
||
for line in text.splitlines():
|
||
cleaned = _clean_answer_line(line)
|
||
if cleaned and not re.match(
|
||
r"(?i)^(?:final answers|explanation|reasoning)\s*:?'?",
|
||
cleaned,
|
||
):
|
||
candidates.append(cleaned)
|
||
if candidates:
|
||
return [candidates[-1]]
|
||
return [""]
|
||
|
||
|
||
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(
|
||
f"Qwen2.5 requires transformers>=4.37.0; found {transformers.__version__}."
|
||
)
|
||
model_config_path = REPO_DIR / "config.json"
|
||
if not model_config_path.exists():
|
||
raise RuntimeError(
|
||
"Qwen2.5-14B-Instruct-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":
|
||
raise RuntimeError(
|
||
"This script expects the official Qwen2.5-14B-Instruct-AWQ snapshot "
|
||
"(model_type=qwen2, quant_method=awq)."
|
||
)
|
||
|
||
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()
|
||
# 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:
|
||
messages = [
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": _problem_prompt(row, fitted_rag, retry_note)},
|
||
]
|
||
rendered = tokenizer.apply_chat_template(
|
||
messages,
|
||
tokenize=False,
|
||
add_generation_prompt=True,
|
||
)
|
||
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 system prompt 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
|
||
|
||
encoded, max_new_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=int(config["max_new_tokens"]),
|
||
)
|
||
encoded = {name: value.to(input_device) for name, value in encoded.items()}
|
||
torch.manual_seed(seed)
|
||
if torch.cuda.is_available():
|
||
torch.cuda.manual_seed_all(seed)
|
||
|
||
# These values match the generation configuration shipped with Qwen2.5
|
||
# Instruct checkpoints while keeping an explicit challenge-time token cap.
|
||
with torch.inference_mode():
|
||
output = model.generate(
|
||
**encoded,
|
||
max_new_tokens=max_new_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,
|
||
)
|
||
prompt_length = int(encoded["input_ids"].shape[-1])
|
||
generated_ids = output[0, prompt_length:].tolist()
|
||
response_text = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
|
||
return parse_model_response("", response_text)
|
||
|
||
|
||
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"])
|
||
)
|
||
failed_results: list[ParseResult] = []
|
||
last_error = ""
|
||
max_attempts = int(config["max_attempts"])
|
||
for attempt in range(max_attempts):
|
||
retry_note = ""
|
||
if attempt:
|
||
retry_note = (
|
||
"FORMAT RETRY. The previous sample was rejected by the answer parser "
|
||
f"because it had {last_error}. Generate a fresh solution. Even if "
|
||
"uncertain, end with exactly FINAL ANSWERS: followed by at least one "
|
||
"non-empty bare answer line and nothing else in that block."
|
||
)
|
||
result = generate_once(
|
||
model=model,
|
||
tokenizer=tokenizer,
|
||
input_device=input_device,
|
||
system_prompt=system_prompt,
|
||
row=row,
|
||
rag_text=rag_text,
|
||
retry_note=retry_note,
|
||
config=config,
|
||
seed=int(config["seed"]) + row_index * max_attempts + attempt,
|
||
)
|
||
if result.valid:
|
||
if attempt:
|
||
print(f" parser accepted retry {attempt + 1}/{max_attempts}", flush=True)
|
||
return result
|
||
last_error = result.error
|
||
failed_results.append(result)
|
||
print(
|
||
f" parser rejected attempt {attempt + 1}/{max_attempts}: {last_error}",
|
||
flush=True,
|
||
)
|
||
|
||
print(" WARNING: all strict retries failed; using a best-effort final line", flush=True)
|
||
last_result = failed_results[-1]
|
||
return ParseResult(
|
||
thinking_trace=last_result.thinking_trace,
|
||
final_text=last_result.final_text,
|
||
answers=_best_effort_answers(
|
||
[failed_result.final_text for failed_result in failed_results]
|
||
),
|
||
valid=False,
|
||
error=f"all {max_attempts} strict parsing attempts failed",
|
||
)
|
||
|
||
|
||
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": "system",
|
||
"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."
|
||
),
|
||
},
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
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()
|
||
explanation = tokenizer.decode(
|
||
generated_ids, skip_special_tokens=True
|
||
).strip()
|
||
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
|
||
|
||
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: last-token logits hook, parser, empty-answer rejection, "
|
||
"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 Qwen"
|
||
)
|
||
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)
|
||
# Do not create an all-empty submission before the model is known to load.
|
||
# Otherwise a missing/incompatible shard or startup OOM is silently scored as
|
||
# zero EM and zero chrF instead of being reported as a runtime failure.
|
||
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()
|