| """Cache local LM scores, calibrate on independent development data, evaluate a fixed policy.""" |
| import argparse |
| import json |
| import math |
| import re |
| import time |
| import unicodedata |
| from pathlib import Path |
|
|
| from vimeml.benchmarks.ajimee import DEVELOPMENT_FORMAT, ROOT, convert_items, json_bytes, sha |
| from vimeml.benchmarks.evaluate_ajimee import load_export, metrics, read_json, rerank, summarize |
|
|
| CACHE_FORMAT = "ime_hybrid_score_cache_v1" |
| POLICY_FORMAT = "ime_hybrid_policy_v1" |
| FORMULA = "azookey_score + lambda * contextual_log_probability_sum" |
| GRID = [0., .05, .1, .2, .3, .5, .75, 1., 1.5, 2., 3.] |
|
|
|
|
| def output_directory(path): |
| if path.exists() and any(path.iterdir()): |
| raise ValueError("Output is not empty; use a new --output to preserve previous results.") |
| path.mkdir(parents=True, exist_ok=True) |
|
|
|
|
| def canonical(text): |
| return "".join(unicodedata.normalize("NFKC", text).split()) |
|
|
|
|
| def overlap_keys(rows): |
| return ({(canonical(row["query"]), canonical(row["left_context"])) for row in rows}, |
| {canonical(row["left_context"] + answer) for row in rows for answer in row["answers"]}) |
|
|
|
|
| def check_independence(development, excluded): |
| dev_queries, dev_texts = overlap_keys(development) |
| queries, texts = overlap_keys(excluded) |
| if dev_queries & queries or dev_texts & texts: |
| raise ValueError("Development examples overlap the excluded AJIMEE set (normalized query/context or context+answer).") |
|
|
|
|
| def prepare_development(input_path, output, excluded, labels_reviewed=False): |
| raw = input_path.read_bytes() |
| converted, mapping, stats = convert_items(json.loads(raw.decode("utf-8-sig"))) |
| keys, _ = overlap_keys(mapping) |
| if len(keys) != len(mapping): |
| raise ValueError("Duplicate development query/context pair.") |
| if not excluded.exists(): |
| raise ValueError("Excluded AJIMEE source is missing; specify --exclude.") |
| exclude_raw = excluded.read_bytes() |
| _, exclude_mapping, _ = convert_items(json.loads(exclude_raw.decode("utf-8-sig"))) |
| check_independence(mapping, exclude_mapping) |
| output_directory(output) |
| files = {"evaluation_items.json": raw, "ajimee-input.json": json_bytes(converted), "case-map.json": json_bytes(mapping)} |
| for name, content in files.items(): |
| (output / name).write_bytes(content) |
| manifest = {"format": DEVELOPMENT_FORMAT, "status": "complete", "role": "development", |
| "source_path": str(input_path.resolve()), "source_sha256": sha(raw), "stats": stats, |
| "labels_reviewed": labels_reviewed, |
| "label_status": "reviewed_by_user_assertion" if labels_reviewed else "provisional_synthetic_labels", |
| "excluded_source_sha256": sha(exclude_raw), |
| "files_sha256": {name: sha(content) for name, content in files.items()}, |
| "policy": "Synthetic development only; all text/references preserved. No candidate/model-score filtering. Reading and label quality require review."} |
| (output / "manifest.json").write_bytes(json_bytes(manifest)) |
| return manifest |
|
|
|
|
| def fingerprint(metadata, provenance): |
| dictionaries = [] |
| for line in provenance["dictionary-versions.txt"].splitlines(): |
| match = re.match(r"^ ([0-9a-f]{40}) (\S+)", line) |
| if not match: |
| raise ValueError("Invalid/unpinned dictionary version record.") |
| dictionaries.append({"commit": match[1], "path": match[2]}) |
| converter = provenance["converter-version.txt"].strip() |
| if not re.fullmatch(r"[0-9a-f]{40}", converter) or not dictionaries: |
| raise ValueError("Missing converter/dictionary commit.") |
| return {"checkpoint_sha256": metadata["checkpoint_sha256"], "tokenizer_sha256": metadata["tokenizer_sha256"], |
| "model_config": metadata["model"], "precision": metadata["precision"], |
| "converter_commit": converter, "dictionaries": sorted(dictionaries, key=lambda row: row["path"]), |
| "n_best": provenance["n_best"]} |
|
|
|
|
| def save_cache(benchmark, output, checkpoint, tokenizer, device, threads): |
| if output.exists() and any(output.iterdir()): |
| raise ValueError("Cache output is not empty; choose a new --output.") |
| manifest, provenance, rows = load_export(benchmark) |
| import torch |
| from vimeml.training.infer import JapaneseLM |
| torch.set_num_threads(threads) |
| start = time.perf_counter() |
| lm = JapaneseLM(checkpoint, tokenizer, device) |
| rerank(lm, rows) |
| output_directory(output) |
| content = "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows).encode("utf-8") |
| (output / "scores.jsonl").write_bytes(content) |
| cache = {"format": CACHE_FORMAT, "status": "complete", "model": lm.metadata, |
| "benchmark_manifest": manifest, "export_provenance": provenance, |
| "files_sha256": {"scores.jsonl": sha(content)}, "fingerprint": fingerprint(lm.metadata, provenance), |
| "elapsed_seconds": time.perf_counter() - start} |
| (output / "manifest.json").write_bytes(json_bytes(cache)) |
| return cache |
|
|
|
|
| def load_cache(directory): |
| |
| path = directory / "manifest.json" |
| if path.exists(): |
| cache = read_json(path) |
| if cache.get("format") != CACHE_FORMAT: |
| raise ValueError("Unknown score cache format.") |
| else: |
| cache = read_json(directory / "metrics.json") |
| if cache.get("format") != "ajimee_ime_metrics_v1": |
| raise ValueError("Unknown evaluation cache format.") |
| if cache.get("status") != "complete" or not cache.get("model"): |
| raise ValueError("Expected complete model scores.") |
| raw = (directory / "scores.jsonl").read_bytes() |
| if sha(raw) != cache["files_sha256"]["scores.jsonl"]: |
| raise ValueError("Cached scores changed.") |
| rows = [json.loads(line) for line in raw.decode("utf-8").splitlines() if line.strip()] |
| if not rows or len(rows) != cache["benchmark_manifest"]["stats"]["cases"] or len({r["id"] for r in rows}) != len(rows): |
| raise ValueError("Invalid score cache sample count/IDs.") |
| for row in rows: |
| texts = [c["text"] for c in row["candidates"]] |
| if texts != row["orders"]["azookey"] or len(set(texts)) != len(texts): |
| raise ValueError("Cached original candidate order changed.") |
| if any(not math.isfinite(c["score"]) for c in row["candidates"]): |
| raise ValueError("Nonfinite AzooKey score.") |
| scores = row.get("lm_scores", {}).get("contextual") |
| if "lm_context_sum" in row["fallbacks"]: |
| if row["orders"]["lm_context_sum"] != texts: |
| raise ValueError("Invalid whole-case fallback.") |
| elif not scores or [c["text"] for c in scores["candidates"]] != texts: |
| raise ValueError("Cached contextual scores do not align with original candidates.") |
| else: |
| for candidate in scores["candidates"]: |
| probabilities = candidate["token_log_probabilities"] |
| if not probabilities or len(probabilities) != candidate["scored_tokens"] or len(candidate["token_ids"]) != len(probabilities): |
| raise ValueError("Invalid cached token trace.") |
| if not all(math.isfinite(v) for v in probabilities) or not math.isclose(sum(probabilities), candidate["log_probability_sum"], abs_tol=1e-6): |
| raise ValueError("Cached LM total differs from token trace.") |
| cache["fingerprint"] = fingerprint(cache["model"], cache["export_provenance"]) |
| cache["cache_directory"] = str(directory.resolve()) |
| return cache, rows |
|
|
|
|
| def valid_weight(value): |
| if not math.isfinite(value) or value < 0: |
| raise ValueError("lambda must be finite and nonnegative.") |
| return value |
|
|
|
|
| def hybrid_order(candidates, scores, weight): |
| """No labels, normalization, length bonus or candidate filtering on this path.""" |
| valid_weight(weight) |
| if weight == 0: |
| return [c["text"] for c in candidates], [c["score"] for c in candidates] |
| if [c["text"] for c in candidates] != [c["text"] for c in scores]: |
| raise ValueError("Hybrid score alignment mismatch.") |
| values = [c["score"] + weight * s["log_probability_sum"] for c, s in zip(candidates, scores)] |
| if not all(math.isfinite(v) for v in values): |
| raise ValueError("Nonfinite combined score.") |
| indices = sorted(range(len(candidates)), key=lambda i: values[i], reverse=True) |
| return [candidates[i]["text"] for i in indices], values |
|
|
|
|
| def add_hybrid(rows, weight, details=True): |
| valid_weight(weight) |
| for row in rows: |
| if "lm_context_sum" in row["fallbacks"] and weight > 0: |
| row["orders"]["hybrid"] = list(row["orders"]["azookey"]) |
| row["fallbacks"]["hybrid"] = row["fallbacks"]["lm_context_sum"] |
| values = None |
| else: |
| scored = row.get("lm_scores", {}).get("contextual", {}).get("candidates", []) |
| row["orders"]["hybrid"], values = hybrid_order(row["candidates"], scored, weight) |
| row["fallbacks"].pop("hybrid", None) |
| if details: |
| row["hybrid_scores_original_order"] = values |
|
|
|
|
| def script_profile(text): |
| return {"kanji": sum('\u3400' <= c <= '\u9fff' for c in text), |
| "hiragana": sum('\u3040' <= c <= '\u309f' for c in text), |
| "katakana": sum('\u30a0' <= c <= '\u30ff' for c in text), |
| "latin": sum(c.isascii() and c.isalpha() for c in text), "digits": sum(c.isdecimal() for c in text)} |
|
|
|
|
| def diagnostics(row): |
| scored = row.get("lm_scores", {}).get("contextual", {}).get("candidates", []) |
| prefixes, fallback = [], [] |
| for i, candidate in enumerate(scored): |
| byte_count = sum(piece.startswith("<0x") for piece in candidate["pieces"]) |
| if byte_count: |
| fallback.append({"text": candidate["text"], "byte_tokens": byte_count}) |
| for other in scored[i + 1:]: |
| a, b = candidate, other |
| if len(a["token_ids"]) > len(b["token_ids"]): |
| a, b = b, a |
| if len(a["token_ids"]) < len(b["token_ids"]) and b["token_ids"][:len(a["token_ids"])] == a["token_ids"]: |
| prefixes.append({"shorter": a["text"], "longer": b["text"]}) |
| original, selected = row["orders"]["azookey"], row["orders"]["hybrid"] |
| before, after = original[0] if original else "", selected[0] if selected else "" |
| return {"strict_token_prefix_pairs": prefixes, "byte_fallback_candidates": fallback, |
| "boundary_retokenized_candidates": sum(c["boundary_retokenized"] for c in scored), |
| "top1_changed": before != after, "before_script_profile": script_profile(before), |
| "after_script_profile": script_profile(after), |
| "script_profile_changed": script_profile(before) != script_profile(after)} |
|
|
|
|
| def calibrate(cache, rows, grid): |
| if cache["benchmark_manifest"].get("role") != "development" or cache["benchmark_manifest"]["format"] != DEVELOPMENT_FORMAT: |
| raise ValueError("Automatic lambda selection requires an independent development cache; AJIMEE cannot be tuned.") |
| if not cache["benchmark_manifest"].get("labels_reviewed"): |
| raise ValueError("Development readings/references need review before calibration; prepare a reviewed version with --labels-reviewed.") |
| if not grid or len(set(grid)) != len(grid): |
| raise ValueError("Expected a nonempty grid of distinct lambda values.") |
| records = [] |
| for weight in grid: |
| valid_weight(weight) |
| add_hybrid(rows, weight, details=False) |
| records.append({"lambda": weight, "metrics": metrics(rows, "hybrid")}) |
| chosen = min(records, key=lambda record: (-record["metrics"]["top1_correct"], |
| record["metrics"]["regressed_vs_azookey"], record["lambda"])) |
| return {"format": POLICY_FORMAT, "status": "development_selected", "formula": FORMULA, |
| "lambda": chosen["lambda"], "fingerprint": cache["fingerprint"], |
| "selection_rule": "Highest development Top-1; ties prefer fewer regressions; remaining ties prefer smaller lambda.", |
| "development_source_sha256": cache["benchmark_manifest"]["source_sha256"], |
| "development_scores_sha256": cache["files_sha256"]["scores.jsonl"], |
| "excluded_source_sha256": cache["benchmark_manifest"]["excluded_source_sha256"], |
| "development_overlap_keys": {"queries": sorted([list(key) for key in overlap_keys(rows)[0]]), |
| "texts": sorted(overlap_keys(rows)[1])}, |
| "grid": records} |
|
|
|
|
| def validate_policy(policy, cache, rows): |
| if policy.get("format") != POLICY_FORMAT or policy.get("status") != "development_selected" or policy.get("formula") != FORMULA: |
| raise ValueError("Invalid selected policy.") |
| valid_weight(policy["lambda"]) |
| if policy["fingerprint"] != cache["fingerprint"]: |
| raise ValueError("Policy model/tokenizer/converter/dictionary/precision/n_best differs from this cache.") |
| if policy["development_source_sha256"] == cache["benchmark_manifest"]["source_sha256"]: |
| raise ValueError("Selected-policy evaluation must use different data from calibration.") |
| queries, texts = overlap_keys(rows) |
| dev = policy["development_overlap_keys"] |
| if queries & {tuple(key) for key in dev["queries"]} or texts & set(dev["texts"]): |
| raise ValueError("Evaluation examples overlap calibration data.") |
| return policy["lambda"] |
|
|
|
|
| def evaluate(cache, rows, weight, output, policy=None): |
| valid_weight(weight) |
| if output.exists() and any(output.iterdir()): |
| raise ValueError("Output is not empty; choose a new --output.") |
| add_hybrid(rows, weight) |
| changed = [] |
| for row in rows: |
| row["hybrid_diagnostics"] = diagnostics(row) |
| if row["hybrid_diagnostics"]["top1_changed"]: |
| before, after = row["orders"]["azookey"][0], row["orders"]["hybrid"][0] |
| changed.append({"id": row["id"], "context": row["left_context"], "query": row["query"], |
| "answers": row["answers"], "before": before, "after": after, |
| "before_correct": before in row["answers"], "after_correct": after in row["answers"], |
| "diagnostics": row["hybrid_diagnostics"]}) |
| report = {"format": "ime_hybrid_evaluation_v1", "status": "complete", "formula": FORMULA, "lambda": weight, |
| "coefficient_source": "independent_development_policy" if policy else "user_specified_not_calibrated", |
| "policy": policy, "score_cache": cache, "metrics": summarize(rows), |
| "diagnostics": {"strict_prefix_cases": sum(bool(r["hybrid_diagnostics"]["strict_token_prefix_pairs"]) for r in rows), |
| "byte_fallback_cases": sum(bool(r["hybrid_diagnostics"]["byte_fallback_candidates"]) for r in rows), |
| "top1_changed_cases": len(changed), |
| "script_profile_changed_cases": sum(r["hybrid_diagnostics"]["script_profile_changed"] for r in rows)}, |
| "policy_notes": "Same real candidate pool; stable ties; raw scores, no EOS/length normalization/extra filters. Whole-case LM fallback. Diagnostics never alter ranking. Strict reference matching; training overlap unknown."} |
| output_directory(output) |
| (output / "scores.jsonl").write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in rows), encoding="utf-8") |
| (output / "changed-cases.json").write_bytes(json_bytes(changed)) |
| report["files_sha256"] = {name: sha((output / name).read_bytes()) for name in ("scores.jsonl", "changed-cases.json")} |
| (output / "metrics.json").write_bytes(json_bytes(report)) |
| lines = ["# AzooKey + LM组合评分", "", f"λ = {weight:g}；系数来源：{report['coefficient_source']}。", "", |
| "同一冻结候选池；并列保持原顺序；超长/无法完整计分时整条回退。表记诊断不是自动语义判分。", "", |
| "| 策略 | 样本数 | Top-1 | Top-5 | MinCER | 纠正 / 改坏 | 回退 |", |
| "| --- | ---: | ---: | ---: | ---: | ---: | ---: |"] |
| for method in ("azookey", "lm_context_sum", "lm_no_context_sum", "hybrid"): |
| m = report["metrics"]["all"][method] |
| lines.append(f"| {method} | {m['cases']} | {m['top1_correct']} ({m['top1_accuracy']:.1%}) | {m['top5_correct']} ({m['top5_accuracy']:.1%}) | " |
| f"{m['mean_min_cer']:.5f} | {m['corrected_vs_azookey']} / {m['regressed_vs_azookey']} | {m['fallback_cases']} |") |
| lines += ["", "完整分组、版本与hash见metrics.json，逐候选组合分数与诊断见scores.jsonl，所有首选变化见changed-cases.json。", ""] |
| (output / "results.md").write_text("\n".join(lines), encoding="utf-8") |
| return report |
|
|
|
|
| def main(argv=None): |
| parser = argparse.ArgumentParser(description=__doc__) |
| sub = parser.add_subparsers(dest="command", required=True) |
| prepare = sub.add_parser("prepare-dev", help="Prepare reviewed DeepSeek JSON for the same Mac export command.") |
| prepare.add_argument("--input", type=Path, required=True) |
| prepare.add_argument("--output", type=Path, default=ROOT / "artifacts/benchmarks/ime-dev-v1") |
| prepare.add_argument("--exclude", type=Path, default=ROOT / "artifacts/benchmarks/ajimee-jwtd-v2-v1/evaluation_items.json") |
| prepare.add_argument("--labels-reviewed", action="store_true", help="Assert readings and acceptable references have been independently reviewed.") |
| score = sub.add_parser("score", help="Compute local LM once for a real candidate export.") |
| score.add_argument("--benchmark", type=Path, required=True) |
| score.add_argument("--output", type=Path, required=True) |
| score.add_argument("--checkpoint", type=Path, default=ROOT / "artifacts/models/tiny-ja-v1/best.pt") |
| score.add_argument("--tokenizer", type=Path, default=ROOT / "artifacts/tokenizers/ja-unigram-16k-v1") |
| score.add_argument("--device", choices=("cpu", "cuda"), default="cpu") |
| score.add_argument("--threads", type=int, default=4) |
| tune = sub.add_parser("tune", help="Select lambda only on reviewed independent development data.") |
| tune.add_argument("--scores", type=Path, required=True) |
| tune.add_argument("--grid", type=float, nargs="+", default=GRID) |
| tune.add_argument("--output", type=Path, required=True) |
| run = sub.add_parser("evaluate", help="Evaluate one fixed lambda or a saved development policy.") |
| run.add_argument("--scores", type=Path, default=ROOT / "outputs/ime-eval/tiny-ja-v1-ajimee") |
| option = run.add_mutually_exclusive_group(required=True) |
| option.add_argument("--lambda", type=float, dest="weight") |
| option.add_argument("--policy", type=Path) |
| run.add_argument("--output", type=Path, required=True) |
| args = parser.parse_args(argv) |
| try: |
| if args.command == "prepare-dev": |
| result = prepare_development(args.input, args.output, args.exclude, args.labels_reviewed) |
| print(json.dumps(result["stats"], indent=2)) |
| print(f"Mac input: {args.output.resolve() / 'ajimee-input.json'}") |
| elif args.command == "score": |
| if args.threads < 1: |
| raise ValueError("threads must be positive.") |
| save_cache(args.benchmark, args.output, args.checkpoint, args.tokenizer, args.device, args.threads) |
| print(f"Score cache: {args.output.resolve()}") |
| else: |
| cache, rows = load_cache(args.scores) |
| if args.command == "tune": |
| result = calibrate(cache, rows, args.grid) |
| output_directory(args.output) |
| (args.output / "policy.json").write_bytes(json_bytes(result)) |
| print(f"Development-selected lambda: {result['lambda']:g}") |
| print(f"Policy: {args.output.resolve() / 'policy.json'}") |
| else: |
| policy = read_json(args.policy) if args.policy else None |
| weight = validate_policy(policy, cache, rows) if policy else args.weight |
| report = evaluate(cache, rows, weight, args.output, policy) |
| print(json.dumps(report["metrics"]["all"]["hybrid"], indent=2)) |
| print(f"Report: {args.output.resolve() / 'results.md'}") |
| except (ValueError, OSError, KeyError, TypeError) as error: |
| parser.exit(2, f"Error: {error}\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
