259 lines
10 KiB
Python
259 lines
10 KiB
Python
import os
|
|
import subprocess
|
|
import sys
|
|
|
|
|
|
def _install_bundled_deps() -> None:
|
|
"""Install transformers from bundled wheels (eval sandbox has no PyPI access)."""
|
|
wheels_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "wheels")
|
|
if not os.path.isdir(wheels_dir):
|
|
return
|
|
subprocess.run(
|
|
[
|
|
sys.executable,
|
|
"-m",
|
|
"pip",
|
|
"install",
|
|
"-q",
|
|
"--no-index",
|
|
f"--find-links={wheels_dir}",
|
|
"transformers==4.56.2",
|
|
],
|
|
check=True,
|
|
)
|
|
|
|
|
|
_install_bundled_deps()
|
|
|
|
import re
|
|
import csv
|
|
import json
|
|
import shutil
|
|
import tempfile
|
|
import torch
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
|
|
# The repo is the working directory at run time, and there is no network.
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
|
MODEL_ID = "."
|
|
MAX_NEW_TOKENS = 2000
|
|
TEMPERATURE = 0.8
|
|
TOP_P = 0.95
|
|
MAX_ATTEMPTS = 5
|
|
|
|
|
|
def load_tokenizer(model_id: str = "."):
|
|
"""Load tokenizer, converting tokenizer.json for older tokenizers if needed."""
|
|
tokenizer_path = os.path.join(model_id, "tokenizer.json")
|
|
with open(tokenizer_path, encoding="utf-8") as handle:
|
|
data = json.load(handle)
|
|
|
|
merges = data.get("model", {}).get("merges", [])
|
|
if not merges or not isinstance(merges[0], list):
|
|
return AutoTokenizer.from_pretrained(model_id)
|
|
|
|
# Older tokenizers expect merge pairs as "a b" strings, not ["a", "b"] lists.
|
|
data["model"]["merges"] = [" ".join(piece) for piece in merges]
|
|
tmpdir = tempfile.mkdtemp()
|
|
for name in ("tokenizer_config.json", "special_tokens_map.json"):
|
|
src = os.path.join(model_id, name)
|
|
if os.path.isfile(src):
|
|
shutil.copy(src, tmpdir)
|
|
with open(os.path.join(tmpdir, "tokenizer.json"), "w", encoding="utf-8") as handle:
|
|
json.dump(data, handle)
|
|
return AutoTokenizer.from_pretrained(tmpdir)
|
|
|
|
|
|
SYSTEM = (
|
|
"You solve International Linguistics Olympiad problems by reasoning from the "
|
|
"data in CONTEXT you are given to solve the problems in QUERY. \n"
|
|
"There are common TASK TYPES that we specify below, but "
|
|
"you may meet a TASK TYPE you have never seen: read the "
|
|
"instruction and the examples, and answer the QUERY in the same form they use.\n\n"
|
|
"Common TASK TYPES and what to return: \n"
|
|
"`translation`: return the translated form only, in the language the task asks for; \n"
|
|
"`fill_blanks`: return only the missing form for each indicated blank "
|
|
"(beware: this could be many different things: a word, a part of a word or a phonetic transcription---pay close attention to what part of the CONTEXT is missing in QUERY); \n"
|
|
"`match_letters`: return only the option letter (for example A, B, C); \n"
|
|
"`text_to_num`: return the number in digits; \n"
|
|
"`num_to_text`: return the number written out in words, in the language asked; \n"
|
|
"any other type: return exactly what the instruction asks for, nothing else. \n\n"
|
|
"As the first part of your answer, reason step by step about (1) the linguistic "
|
|
"rules that can be deduced from the given examples in CONTEXT, and (2) "
|
|
"how to apply them to the given problems in QUERY, and (3) in what format answers need to be returned (words, numbers, phonetic transcriptions, ...). \n"
|
|
"Then write a draft of the final answer. "
|
|
"Subsequently, compare it with the format requirements again, "
|
|
"and verify it's compliant with the deduced rules, and it is complete, i.e. has an answer for each element in QUERY. "
|
|
"If necessary, correct and refine."
|
|
"Finally, write a line that says exactly `FINAL ANSWERS:` "
|
|
"and, below it, write the answers to the items requested in QUERY (not those in CONTEXT),"
|
|
"one answer per line (separated by \n) in the order the items are asked for in the QUERY -- the "
|
|
"bare answer only, no numbering, no quotes, no extra text, according to the given TASK TYPE."
|
|
)
|
|
|
|
tok = load_tokenizer(MODEL_ID)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID, torch_dtype=torch.float16, device_map="auto"
|
|
).eval()
|
|
|
|
with open("/tmp/data/test.csv", encoding="utf-8", newline="") as f:
|
|
test_rows = list(csv.DictReader(f))
|
|
|
|
outputs_queries_types = []
|
|
for r in test_rows:
|
|
|
|
# Create the prompt.
|
|
messages = [
|
|
{"role": "system", "content": SYSTEM},
|
|
{"role": "user", "content":
|
|
f"CONTEXT:{r['context'].strip()}\nTASK TYPE:`{r['task_type']}`\n\nQUERY:{r['query'].strip()}"},
|
|
]
|
|
ids = tok.apply_chat_template(
|
|
messages, add_generation_prompt=True, return_tensors="pt",
|
|
).to(model.device)
|
|
|
|
# Generate the answer.
|
|
done = False
|
|
attempts = 0
|
|
while not done and attempts < MAX_ATTEMPTS:
|
|
with torch.no_grad():
|
|
out = model.generate(
|
|
ids,
|
|
max_new_tokens=MAX_NEW_TOKENS,
|
|
do_sample=True,
|
|
temperature=TEMPERATURE,
|
|
top_p=TOP_P,)
|
|
text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
|
|
# IF NO FINAL ANSWER keyword is used, try again.
|
|
attempts += 1
|
|
if "final answer" not in text.lower() or text.lower().split('final answer')[1].split('\n')==0:
|
|
print(f'TRYING AGAIN...attempts #{attempts+1}/{MAX_ATTEMPTS}')
|
|
else:
|
|
done = True
|
|
|
|
outputs_queries_types.append((text, r['id'], r['query'], r['task_type']))
|
|
print(f"{len(outputs_queries_types)}/{len(test_rows)} done", flush=True)
|
|
|
|
# Postprocess and store the answers.
|
|
def expected_answer_count(query: str, task_type: str) -> int:
|
|
if task_type == "match_letters":
|
|
numbered = re.findall(r"^\s*\d+\.", query, re.MULTILINE)
|
|
return len(numbered) or 1
|
|
|
|
if "blanks" in query.lower():
|
|
range_match = re.search(r"\((\d+)-(\d+)\)", query)
|
|
if range_match:
|
|
return int(range_match.group(2)) - int(range_match.group(1)) + 1
|
|
return len(re.findall(r"\(\d+\)", query)) or 1
|
|
|
|
numbered = re.findall(r"^\s*\d+[.)]", query, re.MULTILINE)
|
|
return len(numbered) or 1
|
|
|
|
|
|
def split_single_line_answer(text: str, expected: int, task_type: str) -> list[str]:
|
|
text = text.strip()
|
|
if expected <= 1:
|
|
return [text]
|
|
|
|
def try_split(pattern: str) -> list[str] | None:
|
|
parts = [part.strip() for part in re.split(pattern, text) if part.strip()]
|
|
return parts if len(parts) == expected else None
|
|
|
|
if task_type == "match_letters":
|
|
for pattern in (r"\s+", r",\s*", r";\s*"):
|
|
if result := try_split(pattern):
|
|
return result
|
|
letters = re.findall(r"[A-Za-z]", text)
|
|
if len(letters) == expected:
|
|
return [letter.upper() for letter in letters]
|
|
return [text]
|
|
|
|
if task_type in ("text_to_num", "num_to_text"):
|
|
for pattern in (r",\s*", r";\s*", r"\s+"):
|
|
if result := try_split(pattern):
|
|
return result
|
|
return [text]
|
|
|
|
for pattern in (r";\s*", r",\s*"):
|
|
if result := try_split(pattern):
|
|
return result
|
|
return [text]
|
|
|
|
|
|
def postprocess_answer(text, query, task_type):
|
|
"""Keep only the lines after the last 'FINAL ANSWERS:' marker, one answer per line,
|
|
stopping at the first empty line. If eval_type is multiple and only one line as answer, split at whitespace."""
|
|
# Updated regex to be more flexible with surrounding characters
|
|
marker_match = list(re.finditer(r"(?im)^[^\w\n]*final answers?[^\w\n]*:?\s*$", text))
|
|
if marker_match:
|
|
text_after_marker = text[marker_match[-1].end():]
|
|
#print('FOUND FINAL ANSWER', text_after_marker)
|
|
else:
|
|
#print("No 'FINAL ANSWERS:' marker found")
|
|
return []
|
|
|
|
answers = []
|
|
answer_lines = text_after_marker.splitlines()
|
|
for i, line in enumerate(answer_lines):
|
|
stripped_line = line.strip('`').strip()
|
|
|
|
# Stop processing if an empty line is encountered (not as first line)
|
|
if stripped_line=='':
|
|
continue
|
|
|
|
# Use a more precise regex to only remove numbering if it's a prefix to other text
|
|
# This ensures that lines which are just numbers (e.g., '1') are not stripped.
|
|
match_numbered_prefix = re.match(r"^\s*\d+[.)]\s+(.*)", stripped_line)
|
|
if match_numbered_prefix:
|
|
cleaned_line = match_numbered_prefix.group(1).strip()
|
|
else:
|
|
cleaned_line = stripped_line
|
|
|
|
# Remove any bold markdown '**'
|
|
cleaned_line = re.sub(r"\*\*", "", cleaned_line).strip()
|
|
|
|
# Specific handling for 'match_letters' task type to strip extra words
|
|
if task_type == 'match_letters':
|
|
parts = [
|
|
part.strip("().[]")
|
|
for part in re.split(r"[\s,;]+", cleaned_line)
|
|
if part.strip()
|
|
]
|
|
if not (len(parts) > 1 and all(re.fullmatch(r"[A-Za-z]", part) for part in parts)):
|
|
match_letter_word = re.match(
|
|
r"^\s*(?:\(([A-Za-z])\)|\[([A-Za-z])\]|([A-Za-z]))\.?:?\s*(.*)$",
|
|
cleaned_line,
|
|
)
|
|
if match_letter_word:
|
|
letter = (
|
|
match_letter_word.group(1)
|
|
or match_letter_word.group(2)
|
|
or match_letter_word.group(3)
|
|
)
|
|
cleaned_line = letter.upper()
|
|
|
|
# Append the cleaned, non-empty line
|
|
if cleaned_line:
|
|
answers.append(cleaned_line)
|
|
|
|
#print('PARSED ANSWERS', answers)
|
|
|
|
# Compare against QUERY length: sometimes model forgets newlines
|
|
#print('QUERY', query)
|
|
expected = expected_answer_count(query, task_type)
|
|
query_len = len(query.splitlines()) - 2
|
|
#print(query_len)
|
|
if len(answers) == 1 and expected > 1:
|
|
answers = split_single_line_answer(answers[0], expected, task_type)
|
|
return answers
|
|
|
|
rows = []
|
|
for answer, row_id, query, task_type in outputs_queries_types:
|
|
answers = postprocess_answer(answer, query, task_type)
|
|
rows.append({"id": row_id, "pred": json.dumps(answers, ensure_ascii=False)})
|
|
with open("submission.csv", "w", encoding="utf-8", newline="") as f:
|
|
writer = csv.DictWriter(f, fieldnames=["id", "pred"])
|
|
writer.writeheader()
|
|
writer.writerows(rows)
|
|
print("wrote submission.csv", flush=True) |