276 lines
11 KiB
Python
276 lines
11 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Aligned 8K HELMET-ICL and OpenAI-MRCR pilot for Qwen checkpoints.
|
||
|
|
|
||
|
|
This is deliberately a small diagnostic runner. HELMET prompt construction
|
||
|
|
matches the ICL schedule used by the local OmniServe reproduction; MRCR uses
|
||
|
|
the official OpenAI rows selected by that reproduction, but renders the
|
||
|
|
messages with the Qwen chat template.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import csv
|
||
|
|
import hashlib
|
||
|
|
import json
|
||
|
|
import math
|
||
|
|
import random
|
||
|
|
import re
|
||
|
|
import string
|
||
|
|
from collections import defaultdict
|
||
|
|
from difflib import SequenceMatcher
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import pandas as pd
|
||
|
|
import torch
|
||
|
|
from datasets import load_dataset
|
||
|
|
from huggingface_hub import hf_hub_download
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
|
|
||
|
|
|
||
|
|
HELMET_SPECS = {
|
||
|
|
"trec_coarse": ("icl_trec_coarse_400shot_balance", 6),
|
||
|
|
"trec_fine": ("icl_trec_fine_400shot_balance", 50),
|
||
|
|
"banking77": ("icl_banking77_360shot_balance", 77),
|
||
|
|
"clinic150": ("icl_clinic150_440shot_balance", 151),
|
||
|
|
"nlu": ("icl_nlu_510shot_balance", 68),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def balanced(data, shots: int, label_field: str, seed: int):
|
||
|
|
rng = random.Random(seed)
|
||
|
|
by_label = defaultdict(list)
|
||
|
|
for item in data:
|
||
|
|
by_label[item[label_field]].append(item)
|
||
|
|
rounds = math.ceil(shots / len(by_label))
|
||
|
|
selected_rounds = [[] for _ in range(rounds)]
|
||
|
|
for examples in by_label.values():
|
||
|
|
indices = rng.sample(range(len(examples)), rounds % len(examples))
|
||
|
|
while len(indices) < rounds:
|
||
|
|
indices += rng.sample(range(len(examples)), min(rounds - len(indices), len(examples)))
|
||
|
|
for index, example_index in enumerate(indices):
|
||
|
|
selected_rounds[index].append(examples[example_index])
|
||
|
|
for examples in selected_rounds:
|
||
|
|
rng.shuffle(examples)
|
||
|
|
return [item for group in selected_rounds for item in group][:shots]
|
||
|
|
|
||
|
|
|
||
|
|
def helmet_source(family: str, seed: int):
|
||
|
|
if family == "trec_coarse":
|
||
|
|
source = load_dataset("CogComp/trec", trust_remote_code=True)
|
||
|
|
return source["train"], source["test"], "text", "coarse_label"
|
||
|
|
if family == "trec_fine":
|
||
|
|
source = load_dataset("CogComp/trec", trust_remote_code=True)
|
||
|
|
return source["train"], source["test"], "text", "fine_label"
|
||
|
|
if family == "banking77":
|
||
|
|
source = load_dataset("PolyAI/banking77", trust_remote_code=True)
|
||
|
|
return source["train"], source["test"], "text", "label"
|
||
|
|
if family == "clinic150":
|
||
|
|
source = load_dataset("clinc/clinc_oos", "plus")
|
||
|
|
return source["train"], source["validation"], "text", "intent"
|
||
|
|
if family == "nlu":
|
||
|
|
source = load_dataset("xingkunliuxtracta/nlu_evaluation_data", trust_remote_code=True)["train"]
|
||
|
|
split = source.train_test_split(test_size=0.1, seed=seed)
|
||
|
|
return split["train"], split["test"], "text", "label"
|
||
|
|
raise ValueError(family)
|
||
|
|
|
||
|
|
|
||
|
|
def build_helmet(seed: int, limit: int):
|
||
|
|
rows = []
|
||
|
|
user_template = (
|
||
|
|
'Use the provided mapping from the text to label to assign a label to the text. '
|
||
|
|
'Only output "label: {{label}}" and nothing else. \n\n{context}\n\n{question}'
|
||
|
|
)
|
||
|
|
for family, (dataset_name, num_labels) in HELMET_SPECS.items():
|
||
|
|
shots = int(dataset_name.split("shot")[0].split("_")[-1])
|
||
|
|
train, test, text_field, label_field = helmet_source(family, seed)
|
||
|
|
samples = balanced(test, limit, label_field, seed)
|
||
|
|
for sample_index, sample in enumerate(samples):
|
||
|
|
local_seed = (int(hashlib.sha256(sample[text_field].encode()).hexdigest(), 16) + seed) % 2**31
|
||
|
|
demos = balanced(train, shots, label_field, local_seed)
|
||
|
|
mapping = list(range(num_labels))
|
||
|
|
random.Random(local_seed).shuffle(mapping)
|
||
|
|
context = "\n\n".join(
|
||
|
|
f"{demo[text_field]}\nlabel: {mapping[int(demo[label_field])]}" for demo in demos
|
||
|
|
)
|
||
|
|
rows.append({
|
||
|
|
"row_id": f"helmet_icl:8k:{family}:{sample_index}",
|
||
|
|
"benchmark": "helmet_icl",
|
||
|
|
"config": family,
|
||
|
|
"prompt": user_template.format(context=context, question=sample[text_field]) + "\nlabel:",
|
||
|
|
"answer": str(mapping[int(sample[label_field])]),
|
||
|
|
"max_new_tokens": 20,
|
||
|
|
"metadata": {"dataset": dataset_name, "shots": shots},
|
||
|
|
})
|
||
|
|
return rows
|
||
|
|
|
||
|
|
|
||
|
|
def build_mrcr(tokenizer, selection_path: Path, limit: int):
|
||
|
|
selection = [json.loads(line) for line in selection_path.read_text().splitlines() if line.strip()]
|
||
|
|
chosen = sorted(
|
||
|
|
(row for row in selection if int(row["context_k"]) == 8),
|
||
|
|
key=lambda row: int(row["sample_index"]),
|
||
|
|
)[:limit]
|
||
|
|
if len(chosen) < limit:
|
||
|
|
raise ValueError(f"MRCR selection contains {len(chosen)} 8K rows, requested {limit}")
|
||
|
|
files = [
|
||
|
|
hf_hub_download("openai/mrcr", filename=f"2needle/2needle_{shard}.parquet", repo_type="dataset")
|
||
|
|
for shard in (0, 1)
|
||
|
|
]
|
||
|
|
frame = pd.concat([pd.read_parquet(path) for path in files], ignore_index=True)
|
||
|
|
rows = []
|
||
|
|
for selected in chosen:
|
||
|
|
source = frame.iloc[int(selected["source_index"])]
|
||
|
|
messages = json.loads(source["prompt"])
|
||
|
|
prompt = tokenizer.apply_chat_template(
|
||
|
|
messages,
|
||
|
|
tokenize=False,
|
||
|
|
add_generation_prompt=True,
|
||
|
|
enable_thinking=False,
|
||
|
|
)
|
||
|
|
answer = str(source["answer"])
|
||
|
|
sample_index = int(selected["sample_index"])
|
||
|
|
rows.append({
|
||
|
|
"row_id": f"mrcr:8k:2needle:{sample_index}",
|
||
|
|
"benchmark": "mrcr",
|
||
|
|
"config": "2needle",
|
||
|
|
"prompt": prompt,
|
||
|
|
"answer": answer,
|
||
|
|
"prefix": str(source["random_string_to_prepend"]),
|
||
|
|
"max_new_tokens": min(768, len(tokenizer(answer, add_special_tokens=False).input_ids) + 64),
|
||
|
|
"metadata": {"source_index": int(selected["source_index"]), "official_tokens": int(selected["official_tokens"])},
|
||
|
|
})
|
||
|
|
return rows
|
||
|
|
|
||
|
|
|
||
|
|
def normalize(text: str) -> str:
|
||
|
|
text = text.lower()
|
||
|
|
text = re.sub(r"\b(a|an|the)\b", " ", text)
|
||
|
|
text = "".join(ch for ch in text if ch not in string.punctuation)
|
||
|
|
return " ".join(text.split())
|
||
|
|
|
||
|
|
|
||
|
|
def score(row, prediction: str) -> float:
|
||
|
|
if row["benchmark"] == "helmet_icl":
|
||
|
|
first = prediction.strip().splitlines()[0] if prediction.strip() else ""
|
||
|
|
parsed = re.sub(r"^label:", "", first, flags=re.I).strip()
|
||
|
|
return float(normalize(parsed) == normalize(row["answer"]))
|
||
|
|
prefix = row["prefix"]
|
||
|
|
if not prediction.startswith(prefix):
|
||
|
|
return 0.0
|
||
|
|
response = prediction.removeprefix(prefix).strip()
|
||
|
|
reference = row["answer"].removeprefix(prefix).strip()
|
||
|
|
return float(SequenceMatcher(None, response, reference).ratio())
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--benchmark", choices=("helmet_icl", "mrcr"), required=True)
|
||
|
|
parser.add_argument("--model-path", required=True)
|
||
|
|
parser.add_argument("--tokenizer-path", required=True)
|
||
|
|
parser.add_argument(
|
||
|
|
"--attn-implementation",
|
||
|
|
default="flash_attention_2",
|
||
|
|
choices=("eager", "sdpa", "flash_attention_2"),
|
||
|
|
help="Transformers attention backend; AHA 4-D local masks require SDPA or eager.",
|
||
|
|
)
|
||
|
|
parser.add_argument("--selection", type=Path)
|
||
|
|
parser.add_argument("--rows", type=Path, help="Frozen aligned prompt rows (JSONL)")
|
||
|
|
parser.add_argument("--method", required=True)
|
||
|
|
parser.add_argument("--output", type=Path, required=True)
|
||
|
|
parser.add_argument("--seed", type=int, default=20260710)
|
||
|
|
parser.add_argument("--limit", type=int, default=1, help="Samples per benchmark config")
|
||
|
|
parser.add_argument(
|
||
|
|
"--stop-new-line",
|
||
|
|
action="store_true",
|
||
|
|
help="Use HELMET's newline stop-token policy during generation.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--duo-sink-size",
|
||
|
|
type=int,
|
||
|
|
help="Optional inference-only override for the AHA sink token count.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--duo-recent-size",
|
||
|
|
type=int,
|
||
|
|
help="Optional inference-only override for the AHA recent-window token count.",
|
||
|
|
)
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, trust_remote_code=True, use_fast=False)
|
||
|
|
if args.rows:
|
||
|
|
rows = [json.loads(line) for line in args.rows.read_text().splitlines() if line.strip()]
|
||
|
|
else:
|
||
|
|
rows = (
|
||
|
|
build_helmet(args.seed, args.limit)
|
||
|
|
if args.benchmark == "helmet_icl"
|
||
|
|
else build_mrcr(tokenizer, args.selection, args.limit)
|
||
|
|
)
|
||
|
|
model = AutoModelForCausalLM.from_pretrained(
|
||
|
|
args.model_path,
|
||
|
|
trust_remote_code=True,
|
||
|
|
torch_dtype=torch.bfloat16,
|
||
|
|
attn_implementation=args.attn_implementation,
|
||
|
|
)
|
||
|
|
for field, value in (
|
||
|
|
("duo_sink_size", args.duo_sink_size),
|
||
|
|
("duo_recent_size", args.duo_recent_size),
|
||
|
|
):
|
||
|
|
if value is not None:
|
||
|
|
if value < 0:
|
||
|
|
parser.error(f"--{field.replace('_', '-')} must be non-negative")
|
||
|
|
previous = getattr(model.config, field, None)
|
||
|
|
setattr(model.config, field, value)
|
||
|
|
print(f"[eval] {field} override: {previous!r} -> {value!r}", flush=True)
|
||
|
|
model = model.to("cuda").eval()
|
||
|
|
stop_token_ids = model.generation_config.eos_token_id
|
||
|
|
stop_token_ids = list(stop_token_ids) if isinstance(stop_token_ids, list) else [stop_token_ids]
|
||
|
|
if args.stop_new_line:
|
||
|
|
newline_tokens = ["\n", "Ċ", "ĊĊ", "<0x0A>"]
|
||
|
|
stop_token_ids += [tokenizer.convert_tokens_to_ids(token) for token in newline_tokens]
|
||
|
|
stop_token_ids = sorted(
|
||
|
|
{
|
||
|
|
token_id
|
||
|
|
for token_id in stop_token_ids
|
||
|
|
if token_id is not None and token_id != tokenizer.unk_token_id
|
||
|
|
}
|
||
|
|
)
|
||
|
|
print(f"[eval] stop_token_ids={stop_token_ids}", flush=True)
|
||
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
outputs = []
|
||
|
|
for ordinal, row in enumerate(rows, 1):
|
||
|
|
encoded = tokenizer(row["prompt"], return_tensors="pt", add_special_tokens=False)
|
||
|
|
prompt_tokens = int(encoded.input_ids.shape[1])
|
||
|
|
print(f"[{args.benchmark}] {ordinal}/{len(rows)} {row['config']} tokens={prompt_tokens}", flush=True)
|
||
|
|
encoded = {key: value.to("cuda") for key, value in encoded.items()}
|
||
|
|
with torch.inference_mode():
|
||
|
|
generated = model.generate(
|
||
|
|
**encoded,
|
||
|
|
do_sample=False,
|
||
|
|
max_new_tokens=row["max_new_tokens"],
|
||
|
|
use_cache=True,
|
||
|
|
eos_token_id=stop_token_ids,
|
||
|
|
pad_token_id=tokenizer.pad_token_id,
|
||
|
|
)
|
||
|
|
prediction = tokenizer.decode(generated[0, prompt_tokens:], skip_special_tokens=True).strip()
|
||
|
|
outputs.append({
|
||
|
|
**row,
|
||
|
|
"method": args.method,
|
||
|
|
"prompt_tokens": prompt_tokens,
|
||
|
|
"prompt_sha256": hashlib.sha256(row["prompt"].encode()).hexdigest(),
|
||
|
|
"prediction": prediction,
|
||
|
|
"score": score(row, prediction),
|
||
|
|
})
|
||
|
|
args.output.write_text("".join(json.dumps(row, ensure_ascii=False) + "\n" for row in outputs))
|
||
|
|
summary = args.output.with_suffix(".summary.csv")
|
||
|
|
with summary.open("w", newline="") as handle:
|
||
|
|
writer = csv.DictWriter(handle, fieldnames=["method", "benchmark", "config", "n", "score", "prompt_tokens"])
|
||
|
|
writer.writeheader()
|
||
|
|
for row in outputs:
|
||
|
|
writer.writerow({"method": args.method, "benchmark": row["benchmark"], "config": row["config"], "n": 1, "score": row["score"], "prompt_tokens": row["prompt_tokens"]})
|
||
|
|
print(summary)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|