File size: 9,708 Bytes
81cffe7 | 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 | """vimeml-tiny-zh-v1 deployment stages; every stage writes a fresh output directory.
Brief: docs/reference/mac-zh-v1-coreml.md. Packages are `artifacts/deployment/tiny-zh-v1-<kind>-v1`;
records go under `outputs/deployment/zh-v1-*`.
baseline Mac PyTorch CPU FP32: fixtures (stateless + forward_cached) and development
rescoring on both tracks, scored and compared with the reference results.
export Inference bundle from deployment.pt.
convert Stateless FP32 package. convert-chunk Chunk-KV FP32 package.
align Stateless fixture gate. align-chunk Chunk-KV gate (fixtures or --stateless).
compress --recipe int8-b32 (deployed) or a survey recipe; --chunk for the chunk-KV package.
evaluate Development, both tracks, through a package (or PyTorch from the bundle), scored
and compared with --reference (a previous evaluate/baseline output).
"""
import argparse
import json
import subprocess
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))
import torch # noqa: E402
from vimeml.benchmarks.zh13_lm_scores import load # noqa: E402
from vimeml.deployment import zh_v1 as zh # noqa: E402
from vimeml.deployment.bundle import environment, fresh_directory, read_json, write_json # noqa: E402
CHECKPOINT = ROOT / "artifacts/models/tiny-zh13-12k-fresh-s1-d384-l9/deployment.pt"
TOKENIZER = ROOT / "artifacts/tokenizers/zh-v13/zh-unigram-12288"
REFERENCE = ROOT / "outputs/deployment/vimeml-tiny-zh-v1-fp32-input"
FROZEN = ROOT / "outputs/benchmarks/ime-standard-zh-v1/frozen-r1"
REFERENCE_SCORES = ROOT / "outputs/benchmarks/zh13-20m/runs/12k-fresh-s1"
REFERENCE_RESULTS = ROOT / "outputs/benchmarks/zh13-20m/judgement/scores"
SCORER = ROOT / "scripts/benchmarks/score_standard_ime_zh.py"
TRACKS = ("full", "abbrev")
EXPECTED_TOP1 = {"full": 1363, "abbrev": 337}
def scorer(*args):
result = subprocess.run([sys.executable, "-I", "-X", "utf8", str(SCORER), *map(str, args)],
check=True, capture_output=True, text=True)
return result.stdout.strip()
def judge(output, label, reference_scores, reference_results, reference_label):
"""score_standard_ime_zh.py score + compare per track; rank differences by item id."""
report = {"label": label, "reference": reference_label, "tracks": {}}
for track in TRACKS:
results = output / "results" / track
line = scorer("score", FROZEN, results, "--track", track, "--scores", output / "scores" / f"{track}.jsonl",
"--name", f"{label}-{track}")
compared = scorer("compare", reference_results(track), results, output / f"compare-{track}.json")
summary = read_json(results / "summary.json")
report["tracks"][track] = {
"scorer": line, "compare": compared, "all": summary["all"],
"rank_differences": zh.rank_differences(reference_scores(track), output / "scores" / f"{track}.jsonl"),
**zh.changed_items(reference_results(track), results)}
print(line, "\n", compared, flush=True)
write_json(output / "judgement.json", report)
return report
@torch.inference_mode()
def baseline(output):
model, processor, meta = load(CHECKPOINT, TOKENIZER / "tokenizer.model")
saved = torch.load(CHECKPOINT, map_location="cpu", weights_only=True)
lm_meta = {"checkpoint_step": saved["step"], "checkpoint_sha256": saved["source_checkpoint_sha256"],
"tokenizer_sha256": meta["tokenizer_model_sha256"]}
info, arrays = zh.read_reference(REFERENCE, type("Meta", (), {"metadata": lm_meta}))
errors = []
for name, case in info["cases"].items():
actual = model(torch.from_numpy(arrays[name + "_input_ids"].copy()).long()).numpy()
valid = case["valid_prefix"]
errors.append({"case": name, "check": "stateless", "valid_prefix": valid,
**zh.error_distribution(arrays[name + "_logits"][:, :valid], actual[:, :valid])})
cache, steps = model.new_cache(), []
ids = torch.from_numpy(arrays["long_input_ids"].copy())
for position in range(ids.shape[1]):
steps.append(model.forward_cached(ids[:, position:position + 1], cache))
errors.append({"case": "long", "check": "forward_cached_steps",
**zh.error_distribution(arrays["long_kv_step_logits"], torch.cat(steps, 1).numpy())})
fresh_directory(output)
passed = all(e["outside_tolerance"] == 0 for e in errors)
write_json(output / "fixtures.json", {"format": "vimeml_zh_v1_mac_pytorch_fixtures_v1", "passed": passed,
"model": {**meta, **lm_meta, "device": "cpu", "threads": torch.get_num_threads()},
"environment": environment(), "reference": info,
"tolerances": {"atol": zh.TOLERANCE, "rtol": zh.TOLERANCE}, "errors": errors})
print(json.dumps({"fixtures_passed": passed, "max_abs": max(e["max_abs"] for e in errors)}), flush=True)
return passed
def main():
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--threads", type=int, default=4)
stages = parser.add_subparsers(dest="stage", required=True)
commands = {name: stages.add_parser(name) for name in (
"baseline", "export", "convert", "convert-chunk", "align", "align-chunk", "compress", "evaluate")}
for name, command in commands.items():
command.add_argument("--output", type=Path, required=True)
if name not in ("baseline", "export", "compress"):
command.add_argument("--bundle", type=Path, required=True)
if name in ("align", "align-chunk", "evaluate"):
command.add_argument("--model", type=Path, required=name == "align-chunk")
commands["align-chunk"].add_argument("--stateless", type=Path,
help="Same-precision stateless package; default compares with PyTorch fixtures")
commands["compress"].add_argument("--source", type=Path, required=True)
commands["compress"].add_argument("--alignment", type=Path, required=True)
commands["compress"].add_argument("--recipe", choices=sorted(zh.RECIPES), default=zh.DEPLOYED_RECIPE)
commands["compress"].add_argument("--chunk", action="store_true")
commands["evaluate"].add_argument("--chunk", action="store_true", help="--model is a chunk-KV package")
commands["evaluate"].add_argument("--reference", type=Path,
help="Earlier evaluate/baseline output; default: the Windows reference run")
commands["evaluate"].add_argument("--reference-label")
args = parser.parse_args()
if args.threads < 1 or args.output.exists():
parser.error("Positive threads and a fresh output path are required.")
torch.set_num_threads(args.threads)
if args.stage == "baseline":
passed = baseline(args.output / "fixtures")
model, processor, meta = load(CHECKPOINT, TOKENIZER / "tokenizer.model")
class Loaded:
def loaded(self):
return model, processor, meta
zh.score_development(Loaded(), FROZEN, args.output / "scores")
judge(args.output, "mac-pytorch", lambda t: REFERENCE_SCORES / f"{t}.jsonl",
lambda t: REFERENCE_RESULTS / f"12k-fresh-s1-{t}", "Windows CUDA FP32 reference")
if not passed:
raise SystemExit("Mac PyTorch fixtures failed; evidence preserved.")
elif args.stage == "export":
zh.export_bundle(CHECKPOINT, TOKENIZER, args.output)
elif args.stage == "convert":
zh.convert(args.bundle, args.output)
elif args.stage == "convert-chunk":
zh.convert_chunk(args.bundle, args.output)
elif args.stage == "compress":
zh.compress(args.source, args.output, args.alignment, args.recipe, args.chunk)
elif args.stage == "align":
lm = zh.CoreMLLMZh(args.bundle, args.model) if args.model else zh.BundleLMZh(args.bundle)
report = zh.align(lm, REFERENCE, args.output)
print(json.dumps({"passed": report["passed"], **report["summary"]}, indent=2))
if not report["passed"]:
raise SystemExit("Strict alignment failed; report preserved. Review ranking quality separately.")
elif args.stage == "align-chunk":
runtime = zh.ChunkRuntimeZh(args.bundle, args.model)
report = zh.align_chunk(runtime, args.bundle, REFERENCE, args.output, args.stateless)
print(json.dumps({"passed": report["passed"], "comparison": report["comparison"], **report["summary"]}, indent=2))
if not report["passed"]:
raise SystemExit("KV chunk alignment failed; evidence preserved.")
elif args.stage == "evaluate":
if args.chunk:
lm = zh.CachedChunkLMZh(args.bundle, args.model)
else:
lm = zh.CoreMLLMZh(args.bundle, args.model) if args.model else zh.BundleLMZh(args.bundle)
zh.score_development(lm, FROZEN, args.output / "scores")
if args.reference:
judge(args.output, lm.metadata.get("package", "pytorch-bundle"),
lambda t: args.reference / "scores" / f"{t}.jsonl", lambda t: args.reference / "results" / t,
args.reference_label or str(args.reference))
else:
judge(args.output, lm.metadata.get("package", "pytorch-bundle"),
lambda t: REFERENCE_SCORES / f"{t}.jsonl", lambda t: REFERENCE_RESULTS / f"12k-fresh-s1-{t}",
"Windows CUDA FP32 reference")
write_json(args.output / "model.json", lm.metadata)
print(f"Output: {args.output.resolve()}")
if __name__ == "__main__":
main()
|