初始化项目，由ModelHub XC社区提供模型
Model: jiamingshan/AHA-L2A-Qwen3-1.7B-repro Source: Original Platform
This commit is contained in:
ModelHub XC
275
recipe/scripts/eval_qwen_external_longctx_pilot.py
Normal file
275
recipe/scripts/eval_qwen_external_longctx_pilot.py
Normal file
@@ -0,0 +1,275 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user