| """V4 tokenizer experiment: shared train-only sample, strict kana checks. |
| |
| Does not open test or blind text. Existing V1/V2 training stays frozen. |
| """ |
| import argparse |
| import hashlib |
| import json |
| import math |
| import random |
| import sys |
| import time |
| import tomllib |
| from collections import Counter |
| from concurrent.futures import ProcessPoolExecutor, as_completed |
| from pathlib import Path |
|
|
| from vimeml.tokenizer.monitor import RunMonitor |
| from vimeml.tokenizer.train import ( |
| ROOT, SPECIAL_IDS, PROBES, aggregate_measurements, dump_json, example, |
| file_sha, measure_range, percentile, train_native, validate, |
| ) |
|
|
| MARKERS = ["<CTX>", "<READ>", "<OUT>"] |
|
|
|
|
| def read_json(path): |
| return json.loads(path.read_text(encoding="utf-8-sig")) |
|
|
|
|
| def permitted_path(path): |
| """Reject held-out roles before any file is opened, including symlinks.""" |
| path = Path(path).resolve() |
| if any("blind" in part.lower() or part.lower() in {"test", "test.txt"} |
| for part in path.parts): |
| raise ValueError(f"Held-out input is forbidden: {path}") |
| return path |
|
|
|
|
| def sample_train(path, output, size, seed, expected, max_bytes, progress): |
| """Explicit reservoir sampling; persist identical bytes for all vocab sizes.""" |
| rng = random.Random(seed) |
| digest = hashlib.sha256() |
| rows = [] |
| count = characters = longest = 0 |
| with permitted_path(path).open("rb") as stream: |
| for count, raw in enumerate(stream, 1): |
| digest.update(raw) |
| text = raw.decode("utf-8").removesuffix("\n").removesuffix("\r") |
| if not text or "▁" in text or any(marker in text for marker in MARKERS): |
| raise ValueError(f"Empty/reserved text in train line {count}") |
| characters += len(text) |
| longest = max(longest, len(text.encode("utf-8"))) |
| if longest > max_bytes: |
| raise ValueError("Train sentence exceeds max_sentence_length; no silent exclusion allowed") |
| if len(rows) < size: |
| rows.append(raw) |
| else: |
| index = rng.randrange(count) |
| if index < size: |
| rows[index] = raw |
| if count % 1_000_000 == 0: |
| progress(count, expected["sentences"]) |
| if count != expected["sentences"] or characters != expected["characters"]: |
| raise ValueError("Train counts differ from frozen corpus metadata") |
| with output.open("wb") as stream: |
| for raw in rows: |
| stream.write(raw if raw.endswith(b"\n") else raw + b"\n") |
| return {"sentences": count, "characters": characters, "max_utf8_bytes": longest, |
| "train_sha256": digest.hexdigest(), "sample_sentences": len(rows), |
| "sample_characters": sum(len(raw.decode("utf-8").rstrip("\r\n")) for raw in rows), |
| "sample_sha256": file_sha(output), "seed": seed, |
| "policy": "Python random.Random(seed) reservoir over every train.txt line; persisted order"} |
|
|
|
|
| def verify_model(model, vocab_size, kana): |
| import sentencepiece as spm |
| from sentencepiece import sentencepiece_model_pb2 as pb |
| processor = spm.SentencePieceProcessor(model_file=str(model)) |
| proto = pb.ModelProto.FromString(model.read_bytes()) |
| if processor.get_piece_size() != vocab_size: |
| raise ValueError("Exact vocabulary size check failed") |
| if any(getattr(processor, f"{name}_id")() != value for name, value in SPECIAL_IDS.items()): |
| raise ValueError("Special IDs must be PAD/UNK/BOS/EOS = 0/1/2/3") |
| normalizer = proto.normalizer_spec |
| if (normalizer.name != "identity" or normalizer.add_dummy_prefix |
| or normalizer.remove_extra_whitespaces or not normalizer.escape_whitespaces): |
| raise ValueError("V4 normalization/whitespace contract failed") |
| if (proto.trainer_spec.model_type != pb.TrainerSpec.UNIGRAM |
| or not proto.trainer_spec.byte_fallback |
| or proto.trainer_spec.required_chars != kana |
| or list(proto.trainer_spec.user_defined_symbols) != MARKERS): |
| raise ValueError("V4 trainer contract failed") |
| for index, marker in enumerate(MARKERS, 4): |
| if (processor.piece_to_id(marker) != index |
| or proto.pieces[index].type != pb.ModelProto.SentencePiece.USER_DEFINED |
| or processor.encode(marker) != [index] |
| or processor.decode([index]) != marker): |
| raise ValueError(f"Reserved marker contract failed: {marker}") |
| checks = [] |
| for char in kana: |
| token = processor.piece_to_id(char) |
| ok = (token != processor.unk_id() and processor.id_to_piece(token) == char |
| and proto.pieces[token].type == pb.ModelProto.SentencePiece.NORMAL |
| and processor.encode(char) == [token] and processor.decode([token]) == char) |
| checks.append({"character": char, "codepoint": f"U+{ord(char):04X}", "id": token, "passed": ok}) |
| if not all(item["passed"] for item in checks): |
| raise ValueError(f"Kana single-character check failed: {[c for c in checks if not c['passed']]}") |
| byte_count = sum(processor.is_byte(i) for i in range(vocab_size)) |
| if byte_count != 256: |
| raise ValueError("Expected exactly 256 byte tokens") |
| probes = [example(processor, text) for text in [*PROBES, *MARKERS, "か\u3099", "ｶﾞ", "ヴヷヸヹヺ"]] |
| if any(not p["roundtrip_ok"] or processor.unk_id() in p["ids"] for p in probes): |
| raise ValueError(f"Probe UNK/roundtrip check failed: {[p for p in probes if not p['roundtrip_ok'] or processor.unk_id() in p['ids']]}") |
| return {"status": "passed", "vocab_size": vocab_size, "special_ids": SPECIAL_IDS, |
| "marker_ids": dict(zip(MARKERS, range(4, 7))), "byte_piece_count": byte_count, |
| "normal_piece_count": vocab_size - 263, "required_kana_count": len(kana), |
| "kana_checks": checks, "probes": probes} |
|
|
|
|
| def measure_corpora(model, paths, expected, workers, batch_size, progress): |
| """Exact statistics on the training sample and full validation, without test.""" |
| import sentencepiece as spm |
| jobs = [] |
| collected = {split: [] for split in paths} |
| for split, path in paths.items(): |
| path = permitted_path(path) |
| size = path.stat().st_size |
| partitions = min(workers * 4, max(1, math.ceil(size / (32 * 1024 * 1024)))) |
| for index in range(partitions): |
| jobs.append((model, path, size * index // partitions, |
| size * (index + 1) // partitions, split, batch_size)) |
| completed = 0 |
| total = sum(item["sentences"] for item in expected.values()) |
| started = time.monotonic() |
| with ProcessPoolExecutor(max_workers=workers) as pool: |
| pending = {pool.submit(measure_range, job): job[4] for job in jobs} |
| for future in as_completed(pending): |
| split = pending[future] |
| result = future.result() |
| collected[split].append(result) |
| completed += result["sentences"] |
| progress(completed, total, time.monotonic() - started) |
| processor = spm.SentencePieceProcessor(model_file=str(model)) |
| stats = {} |
| for split, results in collected.items(): |
| stats[split], _ = aggregate_measurements(processor, results, split) |
| if any(stats[split][key] != expected[split][key] for key in ("sentences", "characters")): |
| raise ValueError(f"Measurement count mismatch: {split}") |
| if stats[split]["unknown_tokens"] or stats[split]["roundtrip_mismatches"]: |
| raise ValueError(f"Measurement UNK/roundtrip failure: {split}") |
| stats[split]["estimated_characters_in_128_content_tokens"] = 128 * stats[split]["characters_per_token"] |
| stats[split]["estimated_characters_in_126_content_tokens_plus_bos_eos"] = 126 * stats[split]["characters_per_token"] |
| return stats |
|
|
|
|
| def text_statistics(processor, texts): |
| lengths = Counter() |
| characters = tokens = byte_tokens = unknown = mismatches = 0 |
| coverage = [] |
| for text in texts: |
| ids = processor.encode(text) |
| lengths[len(ids)] += 1 |
| characters += len(text) |
| tokens += len(ids) |
| byte_tokens += sum(processor.is_byte(i) for i in ids) |
| unknown += ids.count(processor.unk_id()) |
| mismatches += processor.decode(ids) != text |
| |
| proto = processor.encode(text, out_type="immutable_proto") |
| if proto.pieces: |
| coverage.append(proto.pieces[min(126, len(proto.pieces)) - 1].end) |
| else: |
| coverage.append(0) |
| count = len(texts) |
| if unknown or mismatches: |
| raise ValueError("Development text failed UNK/roundtrip verification") |
| return {"texts": count, "characters": characters, "tokens": tokens, |
| "mean_tokens": tokens / count if count else None, |
| "characters_per_token": characters / tokens if tokens else None, |
| "byte_fallback_tokens": byte_tokens, |
| "byte_fallback_token_fraction": byte_tokens / tokens if tokens else None, |
| "unknown_tokens": unknown, "roundtrip_mismatches": mismatches, |
| "token_length_percentiles": {f"p{n}": percentile(lengths, n / 100) for n in (50, 95, 99)}, |
| "fits_128_with_bos_eos_fraction": sum(n for length, n in lengths.items() if length + 2 <= 128) / count if count else None, |
| "mean_characters_retained_in_128_with_bos_eos": sum(coverage) / count if count else None} |
|
|
|
|
| def measure_development(model, directory): |
| """Read the explicit development directory only; no LM or ranking scores.""" |
| import sentencepiece as spm |
| directory = permitted_path(directory) |
| manifest_path = directory / "manifest.json" |
| manifest = read_json(manifest_path) |
| if (directory.name != "development" or manifest.get("split") != "development" |
| or manifest.get("status") != "complete" |
| or manifest.get("format") != "vimeml_standard_ime_candidates_v1"): |
| raise ValueError("Only completed standard development candidates are permitted") |
| from vimeml.benchmarks.standard_ime import load_standard_export |
| _, _, rows = load_standard_export(directory) |
| if not rows or any(row.get("provenance", {}).get("role") != "development" for row in rows): |
| raise ValueError("Benchmark contains a non-development row") |
| processor = spm.SentencePieceProcessor(model_file=str(model)) |
| from vimeml.benchmarks.standard_ime import read |
| originals_path = directory / "evaluation_items.json" |
| originals = read(originals_path) |
| candidates = [candidate["text"] for row in rows for candidate in row["candidates"]] |
| readings = [row["input"] for row in originals] |
| |
| hiragana = ["".join(chr(ord(c) - 0x60) if "ァ" <= c <= "ヶ" else c for c in text) for text in readings] |
| result = {"role": "development", "release_id": manifest["release_id"], "cases": len(rows), |
| "purpose": "Tokenizer statistics only; no probability/ranking/reference accuracy", |
| "candidates": text_statistics(processor, candidates), |
| "reading_original_katakana": text_statistics(processor, readings), |
| "reading_hiragana_diagnostic": text_statistics(processor, hiragana), |
| "context_plus_candidate": text_statistics(processor, [row["left_context"] + c["text"] for row in rows for c in row["candidates"]])} |
| result["input_sha256"] = {str(path): file_sha(path) for path in ( |
| manifest_path, originals_path, directory / "ajimee-results/azookey-candidates.json")} |
| return result |
|
|
|
|
| def run(config_path, dry_run=False): |
| import sentencepiece as spm |
| if spm.__version__ != "0.2.1": |
| raise ValueError("V4 requires sentencepiece==0.2.1") |
| config_path = permitted_path(config_path) |
| config = tomllib.loads(config_path.read_text(encoding="utf-8-sig")) |
| settings = config["training"] |
| validate(settings) |
| if (settings["model_type"] != "unigram" or settings["vocab_size"] != 8192 |
| or settings.get("add_dummy_prefix") is not False |
| or settings.get("user_defined_symbols") != MARKERS |
| or settings["input_sentence_size"] <= 100): |
| raise ValueError("Fixed V4-a tokenizer contract differs from configuration") |
| sizes = config["comparison"]["vocab_sizes"] |
| if sizes != [8192, 12288, 16384]: |
| raise ValueError("Expected 8K/12K/16K comparison") |
| measurement = config["measurement"] |
| if any(type(measurement[k]) is not int or measurement[k] < 1 for k in ("workers", "batch_size")): |
| raise ValueError("Measurement resources must be positive integers") |
| paths = config["paths"] |
| corpus = permitted_path(ROOT / paths["corpus_dir"]) |
| output = permitted_path(ROOT / paths["output_dir"]) |
| kana_path = permitted_path(ROOT / paths["required_chars_file"]) |
| kana = "".join(kana_path.read_text(encoding="utf-8").splitlines()) |
| if not kana or len(set(kana)) != len(kana): |
| raise ValueError("Kana list must be nonempty and unique") |
| benchmark = permitted_path(ROOT / paths["benchmark_development"]) |
| if benchmark.name != "development": |
| raise ValueError("Benchmark path must point specifically to development") |
| if output.exists() and any(output.iterdir()): |
| raise FileExistsError(f"Refusing to overwrite existing experiment: {output}") |
| corpus_manifest = read_json(corpus / "manifest.json") |
| stats = read_json(corpus / "stats.json")["splits"] |
| integrity = read_json(corpus / "integrity-check.json") |
| if (corpus_manifest.get("status") != "complete" or integrity.get("status") != "passed" |
| or file_sha(corpus / "manifest.json") != integrity["manifest_sha256"]): |
| raise ValueError("Frozen corpus integrity metadata failed") |
| inputs = [config_path, kana_path, corpus / "train.txt", corpus / "validation.txt", |
| corpus / "manifest.json", corpus / "stats.json", corpus / "integrity-check.json"] |
| before = {str(p): {"bytes": p.stat().st_size, "mtime_ns": p.stat().st_mtime_ns} for p in inputs} |
| plan = {"output": str(output), "vocab_sizes": sizes, "required_kana_count": len(kana), |
| "trained_on": "train.txt only; one persisted reservoir sample", |
| "sample_size": settings["input_sentence_size"], "seed": settings["seed"], |
| "evaluated_on": ["training sample", "full validation", "standard development"], |
| "test_and_blind_access": False, "lm_training_started": False, |
| "dry_run": dry_run} |
| print(json.dumps(plan, ensure_ascii=False, indent=2), flush=True) |
| if dry_run: |
| return plan |
| output.mkdir(parents=True) |
| dump_json(output / "config.json", config) |
| dump_json(output / "plan.json", plan) |
| sample_path = output / "train-sample.txt" |
| started = time.monotonic() |
| with RunMonitor(output, ROOT / "runs/tokenizer/v4-a-sampling", config["monitoring"]["tensorboard"]) as monitor: |
| monitor.update("checking_inputs") |
| sample = sample_train(corpus / "train.txt", sample_path, settings["input_sentence_size"], |
| settings["seed"], stats["train"], settings["max_sentence_length"], |
| lambda done, total: monitor.update("checking_inputs", preflight_completed_sentences=done, preflight_total_sentences=total)) |
| if sample["train_sha256"] != integrity["exports"]["train.txt"]["sha256"]: |
| raise ValueError("Actual train SHA256 differs from frozen corpus") |
| validation_hash = file_sha(corpus / "validation.txt") |
| if validation_hash != integrity["exports"]["validation.txt"]["sha256"]: |
| raise ValueError("Actual validation SHA256 differs from frozen corpus") |
| dump_json(output / "sampling.json", sample) |
| monitor.update("complete") |
| results = {} |
| for size in sizes: |
| target = output / f"ja-unigram-{size}" |
| target.mkdir() |
| with RunMonitor(target, ROOT / "runs/tokenizer" / f"v4-a-{size}", config["monitoring"]["tensorboard"]) as monitor: |
| monitor.update("training") |
| options = {k: v for k, v in settings.items() if k not in {"seed", "input_sentence_size"}} |
| options.update(input=str(sample_path), model_prefix=str(target / "tokenizer"), vocab_size=size, |
| input_sentence_size=0, shuffle_input_sentence=False, hard_vocab_limit=True, |
| escape_whitespaces=True, required_chars=kana, |
| **{f"{name}_id": value for name, value in SPECIAL_IDS.items()}) |
| train_native(options, settings["seed"], target, monitor) |
| model = target / "tokenizer.model" |
| verification = verify_model(model, size, kana) |
| dump_json(target / "validation.json", verification) |
| monitor.update("measuring") |
| split_stats = measure_corpora(model, {"train_sample": sample_path, "validation": corpus / "validation.txt"}, |
| {"train_sample": {"sentences": sample["sample_sentences"], "characters": sample["sample_characters"]}, "validation": stats["validation"]}, |
| measurement["workers"], measurement["batch_size"], |
| lambda done, total, elapsed: monitor.update("measuring", completed_sentences=done, total_sentences=total, measurement_elapsed_seconds=elapsed)) |
| development = measure_development(model, benchmark) |
| result = {"vocab_size": size, "normal_piece_count": size - 263, |
| "splits": split_stats, "development": development} |
| dump_json(target / "stats.json", result) |
| dump_json(target / "manifest.json", {"status": "complete", "trained_on": "shared train-only sample", |
| "model_sha256": file_sha(model), "vocab_sha256": file_sha(target / "tokenizer.vocab"), |
| "sample_sha256": sample["sample_sha256"], "trainer_options": options, |
| "sentencepiece_version": spm.__version__, "special_ids": SPECIAL_IDS, |
| "marker_ids": verification["marker_ids"], "test_and_blind_access": False}) |
| monitor.final_statistics(split_stats) |
| monitor.update("complete") |
| results[str(size)] = result |
| print(f"[complete] vocab={size}; validation chars/token={split_stats['validation']['characters_per_token']:.4f}", flush=True) |
| after = {str(p): {"bytes": p.stat().st_size, "mtime_ns": p.stat().st_mtime_ns} for p in inputs} |
| if before != after or file_sha(sample_path) != sample["sample_sha256"]: |
| raise ValueError("Inputs changed during training") |
| dump_json(output / "comparison.json", results) |
| dump_json(output / "manifest.json", {"status": "complete", "primary_vocab_size": 8192, |
| "primary_model": str(output / "ja-unigram-8192/tokenizer.model"), "plan": plan, |
| "sample": sample, "validation_sha256": validation_hash, "input_metadata": before, |
| "input_sha256": {str(p): file_sha(p) for p in inputs if p.name not in {"train.txt", "validation.txt"}}, |
| "code_sha256": {str(p): file_sha(p) for p in [Path(__file__), Path(__file__).with_name("train.py"), Path(__file__).with_name("native_train.py"), Path(__file__).with_name("monitor.py")]}, |
| "python_version": sys.version, "elapsed_seconds": time.monotonic() - started, |
| "corpus_ready_for_lm_training": corpus_manifest.get("ready_for_lm_training", False), |
| "reproducibility": "Exact shared sample and artifacts are preserved; native multithreading/platform can affect learned IDs", |
| "context_coverage_note": "Corpus 128-token character estimates are ratio-based, not measured packed-window coverage. Development prefix coverage is measured and includes short texts."}) |
| print(f"[complete] V4 tokenizer experiment: {output}", flush=True) |
| return results |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--config", type=Path, default=ROOT / "configs/tokenizer-v4.toml") |
| parser.add_argument("--dry-run", action="store_true") |
| args = parser.parse_args() |
| run(args.config, args.dry_run) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
