File size: 5,631 Bytes
52ce04b 2f7ee95 52ce04b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | #!/usr/bin/env python3
import os
import sys
import json
import re
import csv
import subprocess
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
# ------------------------------------------------------------------
# 1. Install auto-gptq (PyPI is reachable at runtime)
# ------------------------------------------------------------------
try:
import auto_gptq
except ImportError:
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "auto-gptq"])
import auto_gptq
# ------------------------------------------------------------------
# 2. Offline mode - no internet allowed during evaluation
# ------------------------------------------------------------------
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
# ------------------------------------------------------------------
# 3. Load model from local directory (the repo root)
# ------------------------------------------------------------------
MODEL_DIR = "."
print("Loading tokenizer and model...", flush=True)
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
MODEL_DIR,
device_map="auto",
trust_remote_code=True,
torch_dtype=torch.float16,
)
model.eval()
print("Model loaded.", flush=True)
# ------------------------------------------------------------------
# 4. Helper: count numbered items in query
# ------------------------------------------------------------------
def count_items(query: str) -> int:
numbers = re.findall(r'(?:^|\n)\s*(?:\d+)\.', query)
if not numbers:
lines = query.strip().splitlines()
count = 0
for ln in lines:
if re.match(r'^\s*\d+[\.\)]', ln):
count += 1
return count
return len(numbers)
# ------------------------------------------------------------------
# 5. Parse model output to a JSON list
# ------------------------------------------------------------------
def extract_answers(text: str, expected_count: int) -> list:
json_match = re.search(r'\[\s*".*?"\s*(?:,\s*".*?"\s*)*\]', text, re.DOTALL)
if json_match:
try:
ans = json.loads(json_match.group())
if isinstance(ans, list) and all(isinstance(x, str) for x in ans):
return ans
except:
pass
lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
clean = []
for ln in lines:
m = re.match(r'^\s*\d+[\.\)]\s*(.*)', ln)
if m:
clean.append(m.group(1).strip())
else:
clean.append(ln)
if len(clean) > expected_count:
clean = clean[:expected_count]
elif len(clean) < expected_count:
clean += [""] * (expected_count - len(clean))
return clean
# ------------------------------------------------------------------
# 6. Optional explanation generation (set to True to opt in)
# ------------------------------------------------------------------
EXPLAIN = False # change to True if you want to include explanations
def generate_explanation(context: str, query: str, answer_list: list) -> str:
prompt = (
f"{context}\n\n{query}\n\n"
"My answers are: " + ", ".join(answer_list) + "\n\n"
"Now, provide a concise explanation (max 3 bullet points) of the reasoning behind these answers. "
"Do not repeat the answers, only the reasoning."
)
messages = [
{"role": "system", "content": "You are an expert linguist. Provide a short, human-readable explanation of your reasoning for the given linguistic problem."},
{"role": "user", "content": prompt},
]
ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
with torch.no_grad():
out = model.generate(ids, max_new_tokens=300, do_sample=False)
expl = tokenizer.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
return expl
# ------------------------------------------------------------------
# 7. Main loop over test.csv
# ------------------------------------------------------------------
TEST_CSV = "/tmp/data/test.csv"
df = pd.read_csv(TEST_CSV, dtype=str).fillna("")
rows = []
for idx, r in df.iterrows():
context = r["context"].strip()
query = r["query"].strip()
num_items = count_items(query)
system_msg = (
"You solve International Linguistics Olympiad problems. "
"Answer every numbered item in the query. "
"Output the answers as a JSON list of strings, e.g., [\"answer1\", "answer2\"]. "
"Do not include anything else."
)
user_msg = f"{context}\n\n{query}"
messages = [
{"role": "system", "content": system_msg},
{"role": "user", "content": user_msg},
]
ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model.generate(ids, max_new_tokens=512, do_sample=False)
generated = tokenizer.decode(outputs[0][ids.shape[-1]:], skip_special_tokens=True)
answers = extract_answers(generated, num_items)
print(f"Problem {r['id']}: generated {len(answers)} answers for {num_items} items", flush=True)
explanation = ""
if EXPLAIN:
explanation = generate_explanation(context, query, answers)
rows.append({
"id": r["id"],
"pred": json.dumps(answers, ensure_ascii=False),
"explanation": explanation,
})
pd.DataFrame(rows).to_csv("submission.csv", index=False)
print("wrote submission.csv", flush=True)
|