322 lines
8.4 KiB
Python
322 lines
8.4 KiB
Python
|
|
import gc, torch
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
|
|
import re
|
|
import pandas as pd
|
|
from collections import Counter
|
|
|
|
# Free a model already on the GPU, so re-running this cell doesn't stack a second copy and run out of
|
|
# memory. (If you still hit "out of memory", do Runtime -> Restart session and run this cell just once.)
|
|
model = tok = None
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.reset_peak_memory_stats()
|
|
|
|
MODEL_ID = "."
|
|
|
|
MAX_NEW_TOKENS = 2048
|
|
|
|
generation_config = dict(
|
|
do_sample=True,
|
|
temperature=0.3,
|
|
top_p=0.9,
|
|
top_k=40,
|
|
repetition_penalty=1.05,
|
|
)
|
|
|
|
tok = AutoTokenizer.from_pretrained(MODEL_ID)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID, torch_dtype=torch.float16, device_map="auto",
|
|
).eval()
|
|
print("loaded", MODEL_ID, "| VRAM", round(torch.cuda.max_memory_allocated() / 1e9, 1), "GB")
|
|
|
|
df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
|
|
|
|
SOLVER_SYSTEM = """
|
|
You are an expert International Linguistics Olympiad (IOL) solver.
|
|
|
|
Your task is to infer the hidden linguistic rules ONLY from the provided examples.
|
|
|
|
REQUIRED WORKFLOW:
|
|
1. SCRATCHPAD: List all recurring units (morphemes, words, or sounds) and their meanings.
|
|
2. RULE VERIFICATION: Write down rules for combining these units. Test against examples.
|
|
3. FINAL DERIVATION: Step-by-step derivation for each query item.
|
|
|
|
IMPORTANT RULES:
|
|
- Never rely on outside linguistic knowledge.
|
|
- If the task is 'match_letters', the FINAL ANSWERS must be ONLY the letter (e.g., A, B, C) that corresponds to each query item, one per line. Do NOT output the word itself.
|
|
- If the target language is phonetic (uses brackets [] or special symbols), keep that notation exactly.
|
|
- ABSOLUTELY NO ENGLISH in the FINAL ANSWERS section unless the target language is English.
|
|
- Output exactly one answer per query item.
|
|
|
|
Output format:
|
|
|
|
FINAL ANSWERS:
|
|
[Answer 1]
|
|
[Answer 2]
|
|
... (one per line)
|
|
"""
|
|
|
|
VALIDATOR_SYSTEM = """
|
|
You are an expert IOL solution validator.
|
|
Check if the FINAL ANSWERS match the expected format of the query:
|
|
- For 'match_letters', are they ONLY single letters (A, B, C...)? If they are words, it is INVALID.
|
|
- For 'translation' or 'fill_blanks', are they in the target language (not English)? If there is English, it is INVALID.
|
|
- Is the number of answers correct?
|
|
|
|
Output ONLY 'VALID' or 'INVALID' followed by specific contradictions.
|
|
"""
|
|
|
|
CORRECTOR_SYSTEM = """
|
|
You are correcting an IOL solution.
|
|
- If the task is 'match_letters', replace words with the corresponding labels (A, B, C).
|
|
- Ensure the FINAL ANSWERS contain ONLY the target language forms.
|
|
- Remove all English translations, explanations, or labels like 'Item 1:'.
|
|
|
|
Output exactly:
|
|
FINAL ANSWERS:
|
|
followed by one answer per line.
|
|
"""
|
|
|
|
def parse_answers(text):
|
|
m = list(re.finditer(r"(?im)^\s*FINAL ANSWERS\s*:?\s*$", text))
|
|
if m:
|
|
text = text[m[-1].end():]
|
|
|
|
answers = []
|
|
for line in text.splitlines():
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
|
|
# Clean formatting
|
|
line = re.sub(r"^\d+[.)]\s*", "", line)
|
|
line = re.sub(r"^[-*•]\s*", "", line)
|
|
line = line.strip("'\" ")
|
|
|
|
# Heuristic to filter out leaked reasoning:
|
|
# 1. Skip lines that contain markdown bolding or italics
|
|
if '*' in line or '_' in line:
|
|
continue
|
|
# 2. Skip lines that look like full sentences (too many spaces)
|
|
# unless it's a translation task where the target is a sentence.
|
|
if line.count(' ') > 5 and len(line) > 50:
|
|
continue
|
|
|
|
answers.append(line)
|
|
|
|
return answers
|
|
|
|
def constraint_check(problem, answers):
|
|
|
|
errors = []
|
|
|
|
if len(answers) == 0:
|
|
errors.append("No answers generated.")
|
|
|
|
# No empty answers
|
|
|
|
for i,a in enumerate(answers):
|
|
|
|
if len(a.strip()) == 0:
|
|
errors.append(f"Answer {i+1} is empty.")
|
|
|
|
# Remove duplicate consecutive answers
|
|
|
|
for i in range(1,len(answers)):
|
|
|
|
if answers[i] == answers[i-1]:
|
|
errors.append("Duplicate consecutive answers.")
|
|
|
|
# Very long outputs
|
|
|
|
for a in answers:
|
|
|
|
if len(a) > 120:
|
|
errors.append("Answer too long.")
|
|
|
|
return errors
|
|
|
|
def generate(system_prompt, user_prompt):
|
|
|
|
messages = [
|
|
{"role":"system","content":system_prompt},
|
|
{"role":"user","content":user_prompt},
|
|
]
|
|
|
|
text = tok.apply_chat_template(
|
|
messages,
|
|
tokenize=False,
|
|
add_generation_prompt=True
|
|
)
|
|
|
|
inputs = tok(
|
|
text,
|
|
return_tensors="pt"
|
|
).to(model.device)
|
|
|
|
outputs = model.generate(
|
|
|
|
**inputs,
|
|
|
|
max_new_tokens=MAX_NEW_TOKENS,
|
|
|
|
do_sample=True,
|
|
# Lowered temperature for more stable linguistic reasoning
|
|
temperature=0.1,
|
|
top_p=0.95,
|
|
top_k=50,
|
|
|
|
repetition_penalty=1.05,
|
|
)
|
|
|
|
output = tok.decode(
|
|
outputs[0][inputs.input_ids.shape[1]:],
|
|
skip_special_tokens=True
|
|
)
|
|
|
|
return output
|
|
|
|
def self_consistency(problem, n=3):
|
|
candidates_with_raw = [] # Store (parsed_answers_tuple, raw_text_string)
|
|
|
|
for _ in range(n):
|
|
raw_out_text = generate(
|
|
SOLVER_SYSTEM,
|
|
problem
|
|
)
|
|
parsed_ans = tuple(parse_answers(raw_out_text))
|
|
candidates_with_raw.append((parsed_ans, raw_out_text))
|
|
|
|
# Count occurrences of parsed answers
|
|
parsed_ans_counts = Counter(item[0] for item in candidates_with_raw)
|
|
best_parsed_ans = parsed_ans_counts.most_common(1)[0][0]
|
|
|
|
# Find the raw text that produced the best_parsed_ans (take the first one if multiple)
|
|
best_raw_text = None
|
|
for parsed_ans, raw_text_candidate in candidates_with_raw:
|
|
if parsed_ans == best_parsed_ans:
|
|
best_raw_text = raw_text_candidate
|
|
break
|
|
|
|
return list(best_parsed_ans), best_raw_text
|
|
|
|
def validate(problem, answers):
|
|
|
|
prompt = f"""
|
|
Problem
|
|
|
|
{problem}
|
|
|
|
Candidate solution
|
|
|
|
FINAL ANSWERS:
|
|
|
|
{chr(10).join(answers)}
|
|
"""
|
|
|
|
result = generate(
|
|
VALIDATOR_SYSTEM,
|
|
prompt
|
|
)
|
|
|
|
return result
|
|
|
|
def correct(problem, answers, validator_output):
|
|
|
|
prompt = f"""
|
|
Problem
|
|
|
|
{problem}
|
|
|
|
Previous answer
|
|
|
|
FINAL ANSWERS:
|
|
|
|
{chr(10).join(answers)}
|
|
|
|
Validation
|
|
|
|
{validator_output}
|
|
"""
|
|
|
|
result = generate(
|
|
CORRECTOR_SYSTEM,
|
|
prompt
|
|
)
|
|
|
|
return parse_answers(result)
|
|
|
|
def solve(problem):
|
|
|
|
answers, raw_output_for_best_ans = self_consistency(
|
|
problem,
|
|
n=1
|
|
)
|
|
raw_text = raw_output_for_best_ans
|
|
|
|
for _ in range(3):
|
|
|
|
validator = validate(
|
|
problem,
|
|
answers
|
|
)
|
|
|
|
constraints = constraint_check(
|
|
problem,
|
|
answers
|
|
)
|
|
|
|
if validator.strip() == "VALID" and len(constraints) == 0:
|
|
|
|
return answers, raw_text
|
|
|
|
answers = correct(
|
|
problem,
|
|
answers,
|
|
validator + "\n" + "\n".join(constraints)
|
|
)
|
|
|
|
return answers, raw_text
|
|
results = []
|
|
for i, r in df.iterrows():
|
|
print(f"\n{'=' * 72}\nPROBLEM {i + 1}/{len(df)} -- {r['task_type']}\n{'=' * 72}", flush=True)
|
|
problem_text = f"{r['context'].strip()}\n\n{r['query'].strip()}"
|
|
answers, raw_text = solve(problem_text)
|
|
results.append({"query": r["query"], "raw": raw_text, "pred": answers})
|
|
print(f"\n--> parsed {len(answers)} answers", flush=True)
|
|
SUMMARIZE = (
|
|
"Summarize the following reasoning into a few short bullet points: the rule or pattern found "
|
|
"in the data and the key evidence for the answer. Be concise and structured -- do not repeat "
|
|
"the full reasoning."
|
|
)
|
|
|
|
for i, res in enumerate(results):
|
|
print(f"\n{'=' * 72}\nPROBLEM {i + 1} -- explanation\n{'=' * 72}", flush=True)
|
|
messages = [
|
|
{"role": "system", "content": SUMMARIZE},
|
|
{"role": "user", "content": res["raw"]},
|
|
]
|
|
enc = tok.apply_chat_template(
|
|
messages, add_generation_prompt=True, return_tensors="pt", return_dict=True,
|
|
).to(model.device)
|
|
streamer = TextStreamer(tok, skip_prompt=True, skip_special_tokens=True)
|
|
with torch.no_grad():
|
|
out = model.generate(**enc, max_new_tokens=400, do_sample=False, streamer=streamer)
|
|
res["explanation"] = tok.decode(out[0][enc["input_ids"].shape[-1]:], skip_special_tokens=True).strip()
|
|
print()
|
|
|
|
submission = pd.DataFrame([
|
|
{
|
|
"id": i + 1,
|
|
"pred": res["pred"],
|
|
"explanation": res["explanation"],
|
|
}
|
|
for i, res in enumerate(results)
|
|
])
|
|
|
|
submission.to_csv("submission.csv", index=False)
|
|
|
|
print(submission.head())
|
|
print(f"Saved {len(submission)} predictions to submission.csv")
|