745 lines
27 KiB
Python
745 lines
27 KiB
Python
#!/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()
|