415 lines
15 KiB
Python
415 lines
15 KiB
Python
|
import os
|
||
|
|
import re
|
||
|
|
import sys
|
||
|
|
import json
|
||
|
|
import time
|
||
|
|
import logging
|
||
|
|
|
||
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
||
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||
|
|
|
||
|
|
START_TIME = time.time()
|
||
|
|
SOFT_LIMIT_SECONDS = 26 * 60
|
||
|
|
MODEL_ID = "."
|
||
|
|
|
||
|
|
import pandas as pd
|
||
|
|
import torch
|
||
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||
|
|
|
||
|
|
|
||
|
|
logging.basicConfig(
|
||
|
|
level=logging.INFO,
|
||
|
|
format="%(asctime)s | %(levelname)s | %(message)s",
|
||
|
|
stream=sys.stdout,
|
||
|
|
force=True,
|
||
|
|
)
|
||
|
|
log = logging.getLogger("iol-v1")
|
||
|
|
|
||
|
|
|
||
|
|
SYSTEM = (
|
||
|
|
"You solve International Linguistics Olympiad problems by reasoning only from "
|
||
|
|
"the supplied data. Read the instruction and examples carefully. Preserve every "
|
||
|
|
"requested grammatical feature and output form. For translation, return the full "
|
||
|
|
"translated form; for fill_blanks, return each missing form; for text_to_num, "
|
||
|
|
"return digits; for num_to_text, return the complete written number form. Reason "
|
||
|
|
"systematically and verify the inferred rule against the examples. Then write "
|
||
|
|
"FINAL ANSWERS: exactly once, followed by one bare answer per line in query order. "
|
||
|
|
"Do not number, quote, or explain the final answers. Always answer every item."
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
MATCH_SYSTEM = (
|
||
|
|
"Solve this linguistic matching problem as a complete correspondence. Decompose "
|
||
|
|
"recurring morphemes and compound words before matching. The numbered puzzle-language "
|
||
|
|
"items must be answered in their original order using only option labels. When the "
|
||
|
|
"number of labels equals the number of items, this is a one-to-one assignment: use "
|
||
|
|
"every label exactly once. Resolve high-confidence pairs first, then use the global "
|
||
|
|
"constraint for the remainder. Do not stop early; make a best guess for every item. "
|
||
|
|
"Write FINAL ANSWERS: exactly once, followed by one uppercase label per line. Include "
|
||
|
|
"no numbering, words, meanings, or commentary in the final section."
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
REPAIR_SYSTEM = (
|
||
|
|
"Repair the final output of an International Linguistics Olympiad solution. Use the "
|
||
|
|
"context, query, attempted solution, required count, and valid labels when supplied. "
|
||
|
|
"Return exactly the required number of answers in original item order. For matching, "
|
||
|
|
"return uppercase option labels only and use each label once when explicitly told the "
|
||
|
|
"task is one-to-one. For text_to_num return digits; for num_to_text return complete "
|
||
|
|
"written forms; for fill_blanks include every blank in order. Write FINAL ANSWERS: "
|
||
|
|
"exactly once, then one bare answer per line. No reasoning or commentary."
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
MARKER_RE = re.compile(
|
||
|
|
r"(?im)^\s*(?:[-*•]\s*)?(?:#{1,6}\s*)?(?:\*\*)?"
|
||
|
|
r"final answers?\s*:?(?:\*\*)?\s*$"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def clean_answer_line(line):
|
||
|
|
line = str(line).strip()
|
||
|
|
if re.match(r"^#{1,6}\s+\S", line):
|
||
|
|
return ""
|
||
|
|
line = re.sub(r"^\s*\d+[.)]\s*", "", line)
|
||
|
|
line = re.sub(r"^\s*[-*•]\s+", "", line).strip()
|
||
|
|
if line.startswith("**") and line.endswith("**"):
|
||
|
|
line = line[2:-2].strip()
|
||
|
|
return line
|
||
|
|
|
||
|
|
|
||
|
|
def parse_answers(text):
|
||
|
|
"""Parse both a normal final block and repeated per-item FINAL ANSWERS blocks."""
|
||
|
|
text = str(text)
|
||
|
|
markers = list(MARKER_RE.finditer(text))
|
||
|
|
if not markers:
|
||
|
|
return []
|
||
|
|
|
||
|
|
# Some models write one FINAL ANSWERS marker for every numbered item.
|
||
|
|
if len(markers) > 1:
|
||
|
|
answers = []
|
||
|
|
for index, marker in enumerate(markers):
|
||
|
|
end = markers[index + 1].start() if index + 1 < len(markers) else len(text)
|
||
|
|
for line in text[marker.end():end].splitlines():
|
||
|
|
candidate = clean_answer_line(line)
|
||
|
|
if candidate:
|
||
|
|
answers.append(candidate)
|
||
|
|
break
|
||
|
|
return answers
|
||
|
|
|
||
|
|
answers = []
|
||
|
|
for line in text[markers[0].end():].splitlines():
|
||
|
|
if re.match(r"^\s*#{1,6}\s+\S", line):
|
||
|
|
break
|
||
|
|
candidate = clean_answer_line(line)
|
||
|
|
if candidate:
|
||
|
|
answers.append(candidate)
|
||
|
|
return answers
|
||
|
|
|
||
|
|
|
||
|
|
def numbered_context_items(context):
|
||
|
|
return list(dict.fromkeys(re.findall(r"(?m)^\s*(\d+)[.)]\s+\S", str(context))))
|
||
|
|
|
||
|
|
|
||
|
|
def option_labels(context):
|
||
|
|
"""Extract labels such as A.-R. from a matching context, preserving order."""
|
||
|
|
labels = re.findall(r"(?m)^\s*([A-Z])\s*[.)]\s+\S", str(context))
|
||
|
|
return list(dict.fromkeys(label.upper() for label in labels))
|
||
|
|
|
||
|
|
|
||
|
|
def count_query_items(query, context="", task_type=""):
|
||
|
|
query = str(query).strip()
|
||
|
|
task_type = str(task_type).strip()
|
||
|
|
|
||
|
|
if task_type == "match_letters":
|
||
|
|
items = numbered_context_items(context)
|
||
|
|
if items:
|
||
|
|
return len(items)
|
||
|
|
|
||
|
|
ranges = re.findall(r"\((\d+)\s*[-–—]\s*(\d+)\)", query)
|
||
|
|
if ranges:
|
||
|
|
start, end = map(int, ranges[-1])
|
||
|
|
if end >= start:
|
||
|
|
return end - start + 1
|
||
|
|
|
||
|
|
numbered = re.findall(r"(?m)^\s*\d+[.)]\s+\S", query)
|
||
|
|
if numbered:
|
||
|
|
return len(numbered)
|
||
|
|
|
||
|
|
blanks = re.findall(r"\((\d+)\)", query)
|
||
|
|
if blanks:
|
||
|
|
return len(set(blanks))
|
||
|
|
|
||
|
|
if ":" in query:
|
||
|
|
body = query.split(":", 1)[1].strip()
|
||
|
|
lines = [line.strip() for line in body.splitlines() if line.strip()]
|
||
|
|
if len(lines) > 1:
|
||
|
|
return len(lines)
|
||
|
|
if len(lines) == 1:
|
||
|
|
comma_items = [item.strip() for item in lines[0].split(",") if item.strip()]
|
||
|
|
if len(comma_items) > 1:
|
||
|
|
return len(comma_items)
|
||
|
|
|
||
|
|
return 1
|
||
|
|
|
||
|
|
|
||
|
|
def normalize_matching_answers(answers):
|
||
|
|
normalized = []
|
||
|
|
for answer in answers:
|
||
|
|
match = re.fullmatch(r"\s*([A-Za-z])(?:[.)])?\s*", str(answer))
|
||
|
|
normalized.append(match.group(1).upper() if match else str(answer).strip().upper())
|
||
|
|
return normalized
|
||
|
|
|
||
|
|
|
||
|
|
def matching_is_permutation(expected_n, labels):
|
||
|
|
return len(labels) == expected_n and len(set(labels)) == expected_n
|
||
|
|
|
||
|
|
|
||
|
|
def valid_matching(answers, expected_n, labels):
|
||
|
|
if len(answers) != expected_n:
|
||
|
|
return False
|
||
|
|
if not labels:
|
||
|
|
return all(re.fullmatch(r"[A-Z]", answer or "") for answer in answers)
|
||
|
|
if not all(answer in labels for answer in answers):
|
||
|
|
return False
|
||
|
|
if matching_is_permutation(expected_n, labels):
|
||
|
|
return len(set(answers)) == expected_n and set(answers) == set(labels)
|
||
|
|
return True
|
||
|
|
|
||
|
|
|
||
|
|
def complete_matching_permutation(answers, expected_n, labels):
|
||
|
|
"""Keep valid first occurrences and fill duplicates/missing slots with unused labels."""
|
||
|
|
if not matching_is_permutation(expected_n, labels):
|
||
|
|
return answers
|
||
|
|
completed = []
|
||
|
|
used = set()
|
||
|
|
holes = []
|
||
|
|
for answer in answers[:expected_n]:
|
||
|
|
answer = str(answer).strip().upper()
|
||
|
|
if answer in labels and answer not in used:
|
||
|
|
completed.append(answer)
|
||
|
|
used.add(answer)
|
||
|
|
else:
|
||
|
|
holes.append(len(completed))
|
||
|
|
completed.append(None)
|
||
|
|
while len(completed) < expected_n:
|
||
|
|
holes.append(len(completed))
|
||
|
|
completed.append(None)
|
||
|
|
unused = [label for label in labels if label not in used]
|
||
|
|
for position, label in zip(holes, unused):
|
||
|
|
completed[position] = label
|
||
|
|
return completed
|
||
|
|
|
||
|
|
|
||
|
|
def write_submission(rows):
|
||
|
|
pd.DataFrame(rows, columns=["id", "pred"]).to_csv("submission.csv", index=False)
|
||
|
|
|
||
|
|
|
||
|
|
def generate_text(messages, max_new_tokens):
|
||
|
|
input_ids = tokenizer.apply_chat_template(
|
||
|
|
messages, add_generation_prompt=True, return_tensors="pt"
|
||
|
|
).to(model.device)
|
||
|
|
with torch.no_grad():
|
||
|
|
output = model.generate(
|
||
|
|
input_ids,
|
||
|
|
max_new_tokens=max_new_tokens,
|
||
|
|
do_sample=False,
|
||
|
|
use_cache=True,
|
||
|
|
pad_token_id=tokenizer.eos_token_id,
|
||
|
|
)
|
||
|
|
return tokenizer.decode(
|
||
|
|
output[0][input_ids.shape[-1]:], skip_special_tokens=True
|
||
|
|
).strip()
|
||
|
|
|
||
|
|
|
||
|
|
def validate_submission(rows, source_df):
|
||
|
|
errors = []
|
||
|
|
if len(rows) != len(source_df):
|
||
|
|
errors.append(f"row count {len(rows)} != {len(source_df)}")
|
||
|
|
ids = [str(row["id"]) for row in rows]
|
||
|
|
expected_ids = source_df["id"].astype(str).tolist()
|
||
|
|
if ids != expected_ids:
|
||
|
|
errors.append("IDs or row order differ from test.csv")
|
||
|
|
if len(ids) != len(set(ids)):
|
||
|
|
errors.append("duplicate IDs")
|
||
|
|
for position, row in enumerate(rows):
|
||
|
|
try:
|
||
|
|
predictions = json.loads(row["pred"])
|
||
|
|
except Exception as error:
|
||
|
|
errors.append(f"row {position}: invalid JSON ({error})")
|
||
|
|
continue
|
||
|
|
source = source_df.iloc[position]
|
||
|
|
expected_n = count_query_items(
|
||
|
|
source["query"], source["context"], source.get("task_type", "")
|
||
|
|
)
|
||
|
|
if not isinstance(predictions, list):
|
||
|
|
errors.append(f"row {position}: pred is not a list")
|
||
|
|
elif len(predictions) != expected_n:
|
||
|
|
errors.append(f"row {position}: {len(predictions)} != {expected_n} answers")
|
||
|
|
elif not all(isinstance(answer, str) and answer != "" for answer in predictions):
|
||
|
|
errors.append(f"row {position}: empty or non-string answer")
|
||
|
|
return errors
|
||
|
|
|
||
|
|
|
||
|
|
log.info("Loading local model")
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=True)
|
||
|
|
model = AutoModelForCausalLM.from_pretrained(
|
||
|
|
MODEL_ID,
|
||
|
|
torch_dtype=torch.float16,
|
||
|
|
device_map="auto",
|
||
|
|
local_files_only=True,
|
||
|
|
).eval()
|
||
|
|
log.info("Model loaded in %.1fs", time.time() - START_TIME)
|
||
|
|
|
||
|
|
df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
|
||
|
|
log.info("Loaded %d test rows", len(df))
|
||
|
|
|
||
|
|
# Always keep a structurally complete CSV on disk.
|
||
|
|
submission_rows = []
|
||
|
|
for _, row in df.iterrows():
|
||
|
|
expected_n = count_query_items(
|
||
|
|
row["query"], row["context"], row.get("task_type", "")
|
||
|
|
)
|
||
|
|
labels = option_labels(row["context"]) if row.get("task_type", "") == "match_letters" else []
|
||
|
|
fallback = labels[:expected_n] if len(labels) >= expected_n else ["?"] * expected_n
|
||
|
|
submission_rows.append({
|
||
|
|
"id": row["id"],
|
||
|
|
"pred": json.dumps(fallback, ensure_ascii=False),
|
||
|
|
})
|
||
|
|
write_submission(submission_rows)
|
||
|
|
|
||
|
|
for position, (_, row) in enumerate(df.iterrows()):
|
||
|
|
task_type = row.get("task_type", "")
|
||
|
|
expected_n = count_query_items(row["query"], row["context"], task_type)
|
||
|
|
labels = option_labels(row["context"]) if task_type == "match_letters" else []
|
||
|
|
one_to_one = task_type == "match_letters" and matching_is_permutation(expected_n, labels)
|
||
|
|
elapsed = time.time() - START_TIME
|
||
|
|
remaining = SOFT_LIMIT_SECONDS - elapsed
|
||
|
|
if remaining > 10 * 60:
|
||
|
|
main_budget = 1536
|
||
|
|
elif remaining > 5 * 60:
|
||
|
|
main_budget = 768
|
||
|
|
else:
|
||
|
|
main_budget = 384
|
||
|
|
|
||
|
|
log.info(
|
||
|
|
"row=%d/%d type=%s expected=%d labels=%d bijection=%s elapsed=%.1fs budget=%d",
|
||
|
|
position + 1,
|
||
|
|
len(df),
|
||
|
|
task_type,
|
||
|
|
expected_n,
|
||
|
|
len(labels),
|
||
|
|
one_to_one,
|
||
|
|
elapsed,
|
||
|
|
main_budget,
|
||
|
|
)
|
||
|
|
|
||
|
|
if task_type == "match_letters":
|
||
|
|
user_content = (
|
||
|
|
f"CONTEXT:\n{row['context'].strip()}\n\nQUERY:\n{row['query'].strip()}\n\n"
|
||
|
|
f"Required answer count: {expected_n}.\n"
|
||
|
|
f"Valid labels: {', '.join(labels) if labels else 'infer from context'}."
|
||
|
|
)
|
||
|
|
system_prompt = MATCH_SYSTEM
|
||
|
|
else:
|
||
|
|
user_content = (
|
||
|
|
f"TASK TYPE: {task_type}\n\nCONTEXT:\n{row['context'].strip()}\n\n"
|
||
|
|
f"QUERY:\n{row['query'].strip()}\n\nRequired answer count: {expected_n}."
|
||
|
|
)
|
||
|
|
system_prompt = SYSTEM
|
||
|
|
|
||
|
|
try:
|
||
|
|
raw_text = generate_text(
|
||
|
|
[
|
||
|
|
{"role": "system", "content": system_prompt},
|
||
|
|
{"role": "user", "content": user_content},
|
||
|
|
],
|
||
|
|
max_new_tokens=main_budget,
|
||
|
|
)
|
||
|
|
log.info("row=%d primary_raw=%r", position + 1, raw_text)
|
||
|
|
answers = parse_answers(raw_text)
|
||
|
|
original_count = len(answers)
|
||
|
|
if task_type == "match_letters":
|
||
|
|
answers = normalize_matching_answers(answers)
|
||
|
|
valid = valid_matching(answers, expected_n, labels)
|
||
|
|
else:
|
||
|
|
valid = len(answers) == expected_n and all(str(answer).strip() for answer in answers)
|
||
|
|
|
||
|
|
log.info("row=%d primary parsed=%d valid=%s", position + 1, original_count, valid)
|
||
|
|
|
||
|
|
if not valid and remaining > 3 * 60:
|
||
|
|
repair_content = (
|
||
|
|
f"TASK TYPE: {task_type}\nREQUIRED ANSWER COUNT: {expected_n}\n"
|
||
|
|
f"ONE-TO-ONE MATCHING: {one_to_one}\n"
|
||
|
|
f"VALID LABELS: {', '.join(labels)}\n\n"
|
||
|
|
f"CONTEXT:\n{row['context'].strip()}\n\nQUERY:\n{row['query'].strip()}\n\n"
|
||
|
|
f"ATTEMPTED SOLUTION:\n{raw_text}"
|
||
|
|
)
|
||
|
|
repair_text = generate_text(
|
||
|
|
[
|
||
|
|
{"role": "system", "content": REPAIR_SYSTEM},
|
||
|
|
{"role": "user", "content": repair_content},
|
||
|
|
],
|
||
|
|
max_new_tokens=768 if remaining > 8 * 60 else 384,
|
||
|
|
)
|
||
|
|
log.info("row=%d repair_raw=%r", position + 1, repair_text)
|
||
|
|
repaired = parse_answers(repair_text)
|
||
|
|
if task_type == "match_letters":
|
||
|
|
repaired = normalize_matching_answers(repaired)
|
||
|
|
repaired_valid = valid_matching(repaired, expected_n, labels)
|
||
|
|
else:
|
||
|
|
repaired_valid = len(repaired) == expected_n and all(
|
||
|
|
str(answer).strip() for answer in repaired
|
||
|
|
)
|
||
|
|
log.info(
|
||
|
|
"row=%d repair parsed=%d valid=%s",
|
||
|
|
position + 1,
|
||
|
|
len(repaired),
|
||
|
|
repaired_valid,
|
||
|
|
)
|
||
|
|
if repaired_valid or abs(len(repaired) - expected_n) < abs(len(answers) - expected_n):
|
||
|
|
answers = repaired
|
||
|
|
|
||
|
|
if task_type == "match_letters":
|
||
|
|
answers = normalize_matching_answers(answers)
|
||
|
|
if one_to_one and not valid_matching(answers, expected_n, labels):
|
||
|
|
answers = complete_matching_permutation(answers, expected_n, labels)
|
||
|
|
|
||
|
|
answers = [str(answer).strip() for answer in answers[:expected_n]]
|
||
|
|
if len(answers) < expected_n:
|
||
|
|
answers.extend(["?"] * (expected_n - len(answers)))
|
||
|
|
answers = [answer if answer else "?" for answer in answers]
|
||
|
|
|
||
|
|
except Exception as error:
|
||
|
|
log.exception("row=%d failed: %s", position + 1, error)
|
||
|
|
if task_type == "match_letters" and len(labels) >= expected_n:
|
||
|
|
answers = labels[:expected_n]
|
||
|
|
else:
|
||
|
|
answers = ["?"] * expected_n
|
||
|
|
if torch.cuda.is_available():
|
||
|
|
torch.cuda.empty_cache()
|
||
|
|
|
||
|
|
submission_rows[position] = {
|
||
|
|
"id": row["id"],
|
||
|
|
"pred": json.dumps(answers, ensure_ascii=False),
|
||
|
|
}
|
||
|
|
log.info(
|
||
|
|
"row=%d id=%s final_answers=%s",
|
||
|
|
position + 1,
|
||
|
|
row["id"],
|
||
|
|
json.dumps(answers, ensure_ascii=False),
|
||
|
|
)
|
||
|
|
write_submission(submission_rows)
|
||
|
|
log.info("row=%d completed answers=%d/%d", position + 1, len(answers), expected_n)
|
||
|
|
|
||
|
|
errors = validate_submission(submission_rows, df)
|
||
|
|
if errors:
|
||
|
|
for error in errors:
|
||
|
|
log.error("validation: %s", error)
|
||
|
|
raise RuntimeError("submission.csv validation failed")
|
||
|
|
|
||
|
|
log.info(
|
||
|
|
"submission.csv valid rows=%d total_runtime=%.1fs",
|
||
|
|
len(submission_rows),
|
||
|
|
time.time() - START_TIME,
|
||
|
|
)
|