File size: 10,720 Bytes
99297d7 | 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 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 | """Baselines for vimeml-tiny-zh-v1 on the Standard zh benchmark (docs/zh-CN/benchmarks/baselines-zh-v1-plan.md).
engine: a Rime configuration ranks on its own; the first `--n` whole-input Han candidates of the
engine (rime_candidates.py query, the same rule that built the frozen pools) are the ranking.
lm: an external causal LM reranks the frozen pool with the v1 scoring rule: start token + left
context + candidate tokenized jointly, the sum of FP32 full-vocabulary log-probabilities
from the first token where any candidate differs from the context alone, no EOS.
Both write <out>/result/{items.jsonl,summary.json} in score_standard_ime_zh.py's format, so its
`compare` pairs them with the v1 runs. blind needs a committed plan that names the benchmark.
"""
import importlib.util
import json
import math
from pathlib import Path
import subprocess
import sys
from vimeml.benchmarks.zh13_lm_scores import VERSION, common_prefix, read_jsonl
from vimeml.training.data import file_sha
FROZEN = "outputs/benchmarks/ime-standard-zh-v1/frozen-r1"
SCORER = "scripts/benchmarks/score_standard_ime_zh.py"
RIME = "scripts/benchmarks/rime_candidates.py"
OPTIONAL = "artifacts/external-eval/site-packages"
ENGINES = ("ice", "ice-gram", "luna")
EXTERNAL = {
"qwen2.5-0.5b": {"id": "Qwen/Qwen2.5-0.5B", "revision": "060db6499f32faf8b98477b0a26969ef7d8b9987"},
"qwen3-0.6b-base": {"id": "Qwen/Qwen3-0.6B-Base", "revision": "da87bfb608c14b7cf20ba1ce41287e8de496c0cd"},
}
GRAMMAR_SHA256 = "3457b69f5fe0fc4f5d211c878bbfa2e2361b90639952c9845c8e31e9d047cde6"
def scorer_module(root):
spec = importlib.util.spec_from_file_location("score_standard_ime_zh", root / SCORER)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def gate(root, split, plan):
if split not in ("development", "blind"):
raise ValueError(f"Unknown split {split}")
if split == "blind":
if not plan:
raise ValueError("blind needs --blind-plan, a committed plan naming the benchmark")
return scorer_module(root).check_blind_plan(plan)
return None
def load_items(root, split, track):
return {i["id"]: i for i in read_jsonl(root / FROZEN / "items.jsonl")
if i["split"] == split and (track == "full" or i["keys_abbrev"] is not None)}
def engine(root, variant, split, track, out, plan=None, n=20):
root, out = Path(root), Path(out)
plan_record = gate(root, split, plan)
if variant not in ENGINES:
raise ValueError(f"Unknown engine {variant}")
if variant == "ice-gram" and file_sha(root / "datasets/ime-standard-zh-v1/engine/grammar/wanxiang-lts-zh-hans.gram") != GRAMMAR_SHA256:
raise ValueError("grammar model differs from the registered SHA256")
out.mkdir(parents=True)
items = load_items(root, split, track)
with open(out / "inputs.jsonl", "w", encoding="utf-8", newline="\n") as stream:
for row in read_jsonl(root / FROZEN / f"rime-{track}.jsonl"):
if row["id"] in items:
stream.write(json.dumps(row, ensure_ascii=False) + "\n")
done = subprocess.run([sys.executable, "-I", "-X", "utf8", str(root / RIME), "query", "--variant", variant, "--n", str(n),
str(out / "inputs.jsonl"), str(out / "candidates.jsonl")], capture_output=True, text=True, encoding="utf-8")
if done.returncode:
raise SystemExit(f"rime_candidates.py failed:\n{done.stderr}")
found = {r["id"]: r for r in read_jsonl(out / "candidates.jsonl")}
if set(found) != set(items):
raise SystemExit("engine output does not cover every item")
scorer = scorer_module(root)
overlap = {r["id"] for r in read_jsonl(root / FROZEN / "overlap" / "overlap.jsonl")}
manifest = json.loads((root / FROZEN / "manifest.json").read_text(encoding="utf-8"))
rows = []
for item_id in sorted(items):
item, ranked = items[item_id], found[item_id]["candidates"]
accepted = {item["target"], *item["alternatives"]}
rank = next((k + 1 for k, text in enumerate(ranked) if text in accepted), None)
rows.append({"id": item_id, "source": item["source"], "dialogue": f'{item["source"]}:{item["dialogue_id"]}',
"length": scorer.length_bin(item["target"]), "target": item["target"],
"choice": ranked[0] if ranked else None, "rank": rank, "overlap": item_id in overlap})
summary = {"version": VERSION, "track": track, "split": split, "system": f"rime-{variant}",
"engine": found[next(iter(found))]["engine"], "scores": None, "fallbacks": 0, "blind_plan": plan_record,
"labels": manifest["labels"], **scorer.summarize(rows),
"sensitivity_without_overlap": scorer.rates([r for r in rows if not r["overlap"]]),
"overlap_items": sum(r["overlap"] for r in rows)}
result = out / "result"
result.mkdir()
with open(result / "items.jsonl", "w", encoding="utf-8", newline="\n") as stream:
for row in rows:
stream.write(json.dumps(row, ensure_ascii=False) + "\n")
(result / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
return summary
class External:
def __init__(self, root, key, device):
sys.path.insert(0, str(Path(root) / OPTIONAL))
import torch
from huggingface_hub import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer
spec = EXTERNAL[key]
path = snapshot_download(spec["id"], revision=spec["revision"], local_files_only=True)
self.torch, self.device, self.spec = torch, device, spec
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
self.tokenizer = AutoTokenizer.from_pretrained(path)
self.model = AutoModelForCausalLM.from_pretrained(path, dtype=torch.float32).to(device).eval()
self.start = self.tokenizer.convert_tokens_to_ids("<|endoftext|>")
self.parameters = sum(p.numel() for p in self.model.parameters())
def encode(self, text):
return self.tokenizer.encode(text, add_special_tokens=False)
def score(self, context, candidates):
torch = self.torch
context_ids = [self.start, *self.encode(context)]
sequences = []
for candidate in candidates:
ids = self.encode(context + candidate)
if self.tokenizer.decode(ids, clean_up_tokenization_spaces=False) != context + candidate:
raise ValueError("Context + candidate does not roundtrip with the tokenizer")
sequences.append([self.start, *ids])
common = common_prefix([context_ids, *sequences])
if common < 1 or any(len(s) <= common for s in sequences):
raise ValueError("Every candidate needs scored targets")
width = max(map(len, sequences))
ids = torch.full((len(sequences), width), self.start, dtype=torch.long)
attention = torch.zeros_like(ids)
for row, s in enumerate(sequences):
ids[row, :len(s)] = torch.tensor(s)
attention[row, :len(s)] = 1
rows, positions, targets = [], [], []
for row, s in enumerate(sequences):
for position in range(common - 1, len(s) - 1):
rows.append(row); positions.append(position); targets.append(s[position + 1])
with torch.inference_mode():
hidden = self.model.model(input_ids=ids.to(self.device), attention_mask=attention.to(self.device),
use_cache=False).last_hidden_state
selected = hidden[torch.tensor(rows, device=self.device), torch.tensor(positions, device=self.device)]
logits = self.model.lm_head(selected).float().log_softmax(-1)
values = logits.gather(-1, torch.tensor(targets, device=self.device).unsqueeze(-1)).squeeze(-1).cpu()
scores = [0.0] * len(sequences)
for row, value in zip(rows, values.tolist()):
scores[row] += value
if not all(math.isfinite(v) for v in scores):
raise ValueError("Non-finite candidate score")
return scores, any(s[:len(context_ids)] != context_ids for s in sequences)
def lm(root, key, split, track, out, plan=None, device="cuda"):
root, out = Path(root), Path(out)
gate(root, split, plan)
if key not in EXTERNAL:
raise ValueError(f"Unknown model {key}")
out.mkdir(parents=True)
items = load_items(root, split, track)
pools = {p["id"]: p for p in read_jsonl(root / FROZEN / f"pool-{track}.jsonl") if p["id"] in items}
if set(pools) != set(items):
raise SystemExit(f"pool-{track}.jsonl does not cover every {split} item")
model = External(root, key, device)
fallbacks = retokenized = empty = 0
with open(out / "scores.jsonl", "w", encoding="utf-8", newline="\n") as stream:
for item_id in sorted(items):
texts = [c["text"] for c in pools[item_id]["candidates"]]
scores = []
if not texts:
empty += 1
else:
try:
scores, changed = model.score(items[item_id]["left_context"], texts)
retokenized += changed
except ValueError:
fallbacks += 1 # the benchmark rule: this item keeps the pool order
scores = None
stream.write(json.dumps({"id": item_id, "candidates": texts, "scores": scores}, ensure_ascii=False) + "\n")
meta = {"model": model.spec, "parameters": model.parameters, "dtype": "float32", "device": device, "version": VERSION,
"track": track, "split": split, "items": len(items), "empty_pools": empty, "fallbacks": fallbacks,
"boundary_retokenized_items": retokenized, "scores_sha256": file_sha(out / "scores.jsonl"),
"scoring": "start <|endoftext|> + left context + candidate, joint tokenization, common-prefix suffix log-probability sum, FP32, no EOS"}
(out / "meta.json").write_text(json.dumps(meta, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
command = [sys.executable, "-I", "-X", "utf8", str(root / SCORER), "score", str(root / FROZEN), str(out / "result"),
"--track", track, "--split", split, "--scores", str(out / "scores.jsonl"), "--name", key]
if plan:
command += ["--blind-plan", str(plan)]
done = subprocess.run(command, capture_output=True, text=True, encoding="utf-8")
if done.returncode:
raise SystemExit(f"scorer failed:\n{done.stderr}")
return json.loads((out / "result" / "summary.json").read_text(encoding="utf-8"))
|