| """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 |
| 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")) |
|
|
