315 lines
11 KiB
Python
315 lines
11 KiB
Python
|
import os
|
||
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
||
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||
|
|
|
||
|
|
import re
|
||
|
|
import json
|
||
|
|
import time
|
||
|
|
import sys
|
||
|
|
import subprocess
|
||
|
|
|
||
|
|
try:
|
||
|
|
import bitsandbytes # noqa: F401
|
||
|
|
except ImportError:
|
||
|
|
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "bitsandbytes"], check=True)
|
||
|
|
import bitsandbytes
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import pandas as pd
|
||
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
|
||
|
|
|
||
|
|
MODEL_ID = "."
|
||
|
|
TIME_LIMIT_S = 28 * 60
|
||
|
|
START = time.time()
|
||
|
|
|
||
|
|
# Try 8-bit first (much better reasoning fidelity than 4-bit). Falls back to 4-bit
|
||
|
|
# only if it doesn't fit in memory. Precision matters a lot for symbolic pattern work.
|
||
|
|
USE_8BIT_FIRST = True
|
||
|
|
|
||
|
|
|
||
|
|
def time_left():
|
||
|
|
return TIME_LIMIT_S - (time.time() - START)
|
||
|
|
|
||
|
|
|
||
|
|
def load_model():
|
||
|
|
if USE_8BIT_FIRST:
|
||
|
|
try:
|
||
|
|
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
|
||
|
|
tok = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=True)
|
||
|
|
model = AutoModelForCausalLM.from_pretrained(
|
||
|
|
MODEL_ID,
|
||
|
|
quantization_config=bnb_config,
|
||
|
|
device_map="auto",
|
||
|
|
local_files_only=True,
|
||
|
|
).eval()
|
||
|
|
print("Loaded in 8-bit.", flush=True)
|
||
|
|
return tok, model
|
||
|
|
except Exception as e:
|
||
|
|
print(f"8-bit load failed ({e}), falling back to 4-bit.", flush=True)
|
||
|
|
|
||
|
|
bnb_config = BitsAndBytesConfig(
|
||
|
|
load_in_4bit=True,
|
||
|
|
bnb_4bit_quant_type="nf4",
|
||
|
|
bnb_4bit_compute_dtype=torch.float16,
|
||
|
|
bnb_4bit_use_double_quant=True,
|
||
|
|
)
|
||
|
|
tok = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=True)
|
||
|
|
model = AutoModelForCausalLM.from_pretrained(
|
||
|
|
MODEL_ID,
|
||
|
|
quantization_config=bnb_config,
|
||
|
|
device_map="auto",
|
||
|
|
torch_dtype=torch.float16,
|
||
|
|
local_files_only=True,
|
||
|
|
).eval()
|
||
|
|
return tok, model
|
||
|
|
|
||
|
|
|
||
|
|
def generate(messages, max_new_tokens=900, do_sample=False, temperature=0.7):
|
||
|
|
ids = tok.apply_chat_template(
|
||
|
|
messages, add_generation_prompt=True, return_tensors="pt"
|
||
|
|
).to(model.device)
|
||
|
|
gen_kwargs = dict(
|
||
|
|
max_new_tokens=max_new_tokens,
|
||
|
|
pad_token_id=tok.eos_token_id,
|
||
|
|
)
|
||
|
|
if do_sample:
|
||
|
|
gen_kwargs.update(do_sample=True, temperature=temperature, top_p=0.9)
|
||
|
|
else:
|
||
|
|
gen_kwargs.update(do_sample=False, temperature=None, top_p=None)
|
||
|
|
|
||
|
|
with torch.no_grad():
|
||
|
|
out = model.generate(ids, **gen_kwargs)
|
||
|
|
return tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
|
||
|
|
|
||
|
|
|
||
|
|
def count_items(query: str) -> int:
|
||
|
|
nums = re.findall(r"(?m)^\s*(\d+)\s*[.)]\s", query)
|
||
|
|
if nums:
|
||
|
|
return max(int(n) for n in nums)
|
||
|
|
lines = [l for l in query.splitlines() if l.strip()]
|
||
|
|
return max(1, len(lines))
|
||
|
|
|
||
|
|
|
||
|
|
def extract_json_object(text: str):
|
||
|
|
"""Pull the LAST valid JSON object out of a text blob (final-answer JSON should
|
||
|
|
come after the reasoning, so we scan from the end)."""
|
||
|
|
text = text.strip()
|
||
|
|
try:
|
||
|
|
return json.loads(text)
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
# find all {...} candidates, try from the last one backwards
|
||
|
|
candidates = list(re.finditer(r"\{.*?\}", text, re.DOTALL))
|
||
|
|
for m in reversed(candidates):
|
||
|
|
try:
|
||
|
|
return json.loads(m.group(0))
|
||
|
|
except Exception:
|
||
|
|
continue
|
||
|
|
# try greedy full-span match as last resort
|
||
|
|
m = re.search(r"\{.*\}", text, re.DOTALL)
|
||
|
|
if m:
|
||
|
|
try:
|
||
|
|
return json.loads(m.group(0))
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def extract_answers_fallback(text: str, n: int):
|
||
|
|
lines = [l.strip() for l in text.splitlines() if l.strip()]
|
||
|
|
cleaned = []
|
||
|
|
for l in lines:
|
||
|
|
l = re.sub(r"^\s*\d+[.)]\s*", "", l)
|
||
|
|
l = l.strip(" -\t")
|
||
|
|
if l:
|
||
|
|
cleaned.append(l)
|
||
|
|
if len(cleaned) >= n:
|
||
|
|
return cleaned[-n:] # prefer the tail: likely the final-answer block
|
||
|
|
parts = [p.strip() for p in text.replace("\n", ",").split(",") if p.strip()]
|
||
|
|
if len(parts) >= n:
|
||
|
|
return parts[-n:]
|
||
|
|
cleaned = cleaned or parts or [text.strip()]
|
||
|
|
while len(cleaned) < n:
|
||
|
|
cleaned.append(cleaned[-1] if cleaned else "")
|
||
|
|
return cleaned[:n]
|
||
|
|
|
||
|
|
|
||
|
|
def fit_answers(answers, n):
|
||
|
|
answers = list(answers)
|
||
|
|
if len(answers) < n:
|
||
|
|
answers = answers + [answers[-1] if answers else ""] * (n - len(answers))
|
||
|
|
return answers[:n]
|
||
|
|
|
||
|
|
|
||
|
|
# Single combined system prompt: reasoning + mandatory self-verification + strict
|
||
|
|
# final JSON block. This replaces the old two-call reasoning->extraction pipeline,
|
||
|
|
# roughly halving generation cost per row while forcing the model to check its own
|
||
|
|
# rule against every given example (the single highest-leverage habit for IOL-style
|
||
|
|
# problems) before committing to an answer.
|
||
|
|
SOLVE_SYSTEM = (
|
||
|
|
"You are an expert at International Linguistics Olympiad (IOL) problems. You will "
|
||
|
|
"be given example data (context) in an unfamiliar language and a query asking you "
|
||
|
|
"to translate, match, fill blanks, or convert forms.\n\n"
|
||
|
|
"Follow this exact procedure:\n"
|
||
|
|
"1. REASONING: Identify morphemes, word order, and correspondences using ONLY the "
|
||
|
|
"given examples. Do not assume rules from natural languages you know unless the "
|
||
|
|
"data supports them.\n"
|
||
|
|
"2. VERIFY: Explicitly re-derive every example in the context using your proposed "
|
||
|
|
"rules and confirm each one matches. If any mismatch, revise your rule before "
|
||
|
|
"continuing.\n"
|
||
|
|
"3. ANSWER: Apply the verified rule to the query.\n"
|
||
|
|
"4. Output a final line starting with FINAL_JSON: followed by a single JSON object "
|
||
|
|
"with exactly two keys:\n"
|
||
|
|
' "answers": a JSON list of strings, one per numbered item in the query, in the '
|
||
|
|
"same order as the query items.\n"
|
||
|
|
' "explanation": 3-6 short bullet points (use \\n between them) summarizing the '
|
||
|
|
"key rules used.\n"
|
||
|
|
"The FINAL_JSON line must be the last thing you output, and must contain nothing "
|
||
|
|
"else on that line besides FINAL_JSON: {...}."
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def solve_once(problem_text: str, n_items: int, max_new_tokens: int, do_sample=False, temperature=0.7):
|
||
|
|
raw = generate(
|
||
|
|
[
|
||
|
|
{"role": "system", "content": SOLVE_SYSTEM},
|
||
|
|
{"role": "user", "content": problem_text + f"\n\n(The query has {n_items} numbered item(s).)"},
|
||
|
|
],
|
||
|
|
max_new_tokens=max_new_tokens,
|
||
|
|
do_sample=do_sample,
|
||
|
|
temperature=temperature,
|
||
|
|
)
|
||
|
|
|
||
|
|
obj = None
|
||
|
|
m = re.search(r"FINAL_JSON:\s*(\{.*\})", raw, re.DOTALL)
|
||
|
|
if m:
|
||
|
|
try:
|
||
|
|
obj = json.loads(m.group(1))
|
||
|
|
except Exception:
|
||
|
|
obj = extract_json_object(m.group(1))
|
||
|
|
if obj is None:
|
||
|
|
obj = extract_json_object(raw)
|
||
|
|
|
||
|
|
answers, explanation = None, ""
|
||
|
|
if isinstance(obj, dict):
|
||
|
|
a = obj.get("answers")
|
||
|
|
if isinstance(a, list):
|
||
|
|
answers = [str(x) for x in a]
|
||
|
|
explanation = str(obj.get("explanation", "")).strip()
|
||
|
|
|
||
|
|
if answers is None:
|
||
|
|
answers = extract_answers_fallback(raw, n_items)
|
||
|
|
if not explanation:
|
||
|
|
snippet = raw.strip().replace("\n", " ")
|
||
|
|
explanation = (snippet[:400] + "...") if len(snippet) > 400 else snippet
|
||
|
|
|
||
|
|
return fit_answers(answers, n_items), explanation, raw
|
||
|
|
|
||
|
|
|
||
|
|
def majority_vote(all_answers, n_items):
|
||
|
|
"""Per-position majority vote across N sampled runs."""
|
||
|
|
final = []
|
||
|
|
for i in range(n_items):
|
||
|
|
votes = {}
|
||
|
|
for ans_list in all_answers:
|
||
|
|
v = ans_list[i] if i < len(ans_list) else ""
|
||
|
|
votes[v] = votes.get(v, 0) + 1
|
||
|
|
final.append(max(votes.items(), key=lambda kv: kv[1])[0])
|
||
|
|
return final
|
||
|
|
|
||
|
|
|
||
|
|
def solve_row(context: str, query: str, n_items: int, row_budget_s: float):
|
||
|
|
problem_text = f"{context.strip()}\n\n{query.strip()}"
|
||
|
|
|
||
|
|
# Scale tokens to problem size; cap so one bad row can't eat the whole budget.
|
||
|
|
approx_len = len(context) + len(query)
|
||
|
|
max_new = 700 if approx_len < 1500 else 1000
|
||
|
|
max_new = min(max_new, 1100)
|
||
|
|
|
||
|
|
row_start = time.time()
|
||
|
|
|
||
|
|
def row_time_left():
|
||
|
|
return row_budget_s - (time.time() - row_start)
|
||
|
|
|
||
|
|
# First (greedy, most reliable) attempt.
|
||
|
|
answers, explanation, _ = solve_once(problem_text, n_items, max_new, do_sample=False)
|
||
|
|
|
||
|
|
# If time allows, take 2 more sampled attempts and majority-vote for robustness.
|
||
|
|
# This costs roughly the same as the old broken two-pass pipeline, but spends the
|
||
|
|
# extra budget on getting the ANSWER right instead of just reformatting it.
|
||
|
|
if row_time_left() > max_new * 0.03: # crude per-token-time heuristic guard
|
||
|
|
try:
|
||
|
|
extra_runs = [answers]
|
||
|
|
for t in (0.5, 0.8):
|
||
|
|
if row_time_left() < 20:
|
||
|
|
break
|
||
|
|
a2, _, _ = solve_once(problem_text, n_items, max_new, do_sample=True, temperature=t)
|
||
|
|
extra_runs.append(a2)
|
||
|
|
if len(extra_runs) > 1:
|
||
|
|
answers = majority_vote(extra_runs, n_items)
|
||
|
|
except Exception:
|
||
|
|
pass # keep the greedy answer if sampling fails
|
||
|
|
|
||
|
|
return answers, explanation
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
global tok, model
|
||
|
|
tok, model = load_model()
|
||
|
|
|
||
|
|
df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
|
||
|
|
n_rows = len(df)
|
||
|
|
out_rows = []
|
||
|
|
|
||
|
|
for i, r in df.iterrows():
|
||
|
|
rid = r["id"]
|
||
|
|
context, query = r.get("context", ""), r.get("query", "")
|
||
|
|
n_items = count_items(query)
|
||
|
|
|
||
|
|
rows_remaining = n_rows - i
|
||
|
|
# Reserve a small safety margin (30s) for writing the CSV at the end.
|
||
|
|
safe_time_left = max(time_left() - 30, 0)
|
||
|
|
row_budget = safe_time_left / max(rows_remaining, 1)
|
||
|
|
|
||
|
|
if safe_time_left < 20:
|
||
|
|
# Truly out of time: still try a fast, cheap single greedy pass rather
|
||
|
|
# than silently submitting blanks — a bad guess beats an empty string
|
||
|
|
# on chrF and exact-match both.
|
||
|
|
try:
|
||
|
|
answers, explanation, _ = solve_once(
|
||
|
|
f"{context.strip()}\n\n{query.strip()}", n_items, max_new_tokens=300
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
answers = [""] * n_items
|
||
|
|
explanation = f"(error: {e})"
|
||
|
|
else:
|
||
|
|
try:
|
||
|
|
answers, explanation = solve_row(context, query, n_items, row_budget)
|
||
|
|
except Exception as e:
|
||
|
|
answers = [""] * n_items
|
||
|
|
explanation = f"(error: {e})"
|
||
|
|
|
||
|
|
out_rows.append(
|
||
|
|
{
|
||
|
|
"id": rid,
|
||
|
|
"pred": json.dumps(answers, ensure_ascii=False),
|
||
|
|
"explanation": explanation,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
print(
|
||
|
|
f"[{i + 1}/{n_rows}] id={rid} n_items={n_items} "
|
||
|
|
f"row_budget={row_budget:.0f}s time_left={time_left():.0f}s",
|
||
|
|
flush=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Checkpoint after every row so a hard timeout still leaves a submission
|
||
|
|
# instead of losing everything if the process gets killed mid-loop.
|
||
|
|
pd.DataFrame(out_rows, columns=["id", "pred", "explanation"]).to_csv(
|
||
|
|
"submission.csv", index=False
|
||
|
|
)
|
||
|
|
|
||
|
|
print("wrote submission.csv", flush=True)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|