| """V4-only TXT encoding with provenance, strict checks and resumable parts. |
| |
| Original training JSONL was pruned locally. Verified V2 source IDs are reused; |
| all token sequences and row byte positions are regenerated from the frozen TXT. |
| No test/blind text is read, and V1/V2 implementations remain unchanged. |
| """ |
| import argparse |
| import hashlib |
| import json |
| import time |
| import tomllib |
| from collections import Counter |
| from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed |
| from pathlib import Path |
|
|
| from vimeml.tokenizer.encode import dump_json, file_sha, write_numbers |
| from vimeml.tokenizer.encode_parallel import cached_part, output_lock, part_paths |
| from vimeml.tokenizer.monitor import RunMonitor |
| from vimeml.tokenizer.v4 import MARKERS, ROOT, permitted_path, read_json, verify_model |
|
|
| SPLITS = ("train", "validation") |
| FORMAT = "vimeml_sentence_tokens_v1" |
|
|
|
|
| def recorded_hash(manifest, name, parent=None): |
| values = [v for k, v in manifest["input_sha256"].items() |
| if k.replace("\\", "/").rsplit("/", 1)[-1] == name |
| and (parent is None or k.replace("\\", "/").endswith(f"/{parent}/{name}"))] |
| if len(values) != 1: |
| raise ValueError(f"Missing unique recorded corpus hash: {name}") |
| return values[0] |
|
|
|
|
| def scan_txt(path, split, chunk_bytes, expected, progress): |
| """Hash and partition at exact line boundaries, retaining source row indices.""" |
| path = permitted_path(path) |
| digest = hashlib.sha256() |
| jobs = [] |
| count = characters = start = first_row = end = longest = 0 |
| with path.open("rb") as stream: |
| for raw in stream: |
| 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"Invalid/reserved text: {split} row {count}") |
| count += 1 |
| characters += len(text) |
| longest = max(longest, len(text.encode("utf-8"))) |
| end = stream.tell() |
| if end - start >= chunk_bytes: |
| jobs.append({"split": split, "input": str(path), "start": start, "end": end, |
| "first_row": first_row, "sentences": count - first_row}) |
| start, first_row = end, count |
| if count % 1_000_000 == 0: |
| progress(count, expected["sentences"]) |
| if end > start: |
| jobs.append({"split": split, "input": str(path), "start": start, "end": end, |
| "first_row": first_row, "sentences": count - first_row}) |
| if count != expected["sentences"] or characters != expected["characters"]: |
| raise ValueError(f"Corpus count mismatch: {split}") |
| return jobs, {"sentences": count, "characters": characters, "max_utf8_bytes": longest, |
| "sha256": digest.hexdigest()} |
|
|
|
|
| def encode_txt_part(job): |
| import sentencepiece as spm |
| processor = spm.SentencePieceProcessor(model_file=job["model"]) |
| files = part_paths(job) |
| files["done.json"].unlink(missing_ok=True) |
| reserved = set(range(7)) |
| byte_ids = {i for i in range(8192) if processor.is_byte(i)} |
| reverse_sources = {value: name for name, value in job["source_ids"].items()} |
| count = content = stored = characters = byte_tokens = longest = 0 |
| lengths, sources = Counter(), Counter() |
| text_digest, source_digest = hashlib.sha256(), hashlib.sha256() |
| with Path(job["input"]).open("rb") as text_stream, Path(job["source_input"]).open("rb") as source_stream, \ |
| files["tokens.bin"].open("wb") as token_stream, files["offsets.bin"].open("wb") as offset_stream, \ |
| files["rows.bin"].open("wb") as row_stream, files["sources.bin"].open("wb") as out_sources: |
| text_stream.seek(job["start"]) |
| source_stream.seek(job["first_row"]) |
| write_numbers(offset_stream, "Q", [0]) |
| while text_stream.tell() < job["end"]: |
| texts, positions = [], [] |
| for _ in range(job["batch_size"]): |
| position = text_stream.tell() |
| if position >= job["end"]: |
| break |
| raw = text_stream.readline() |
| if not raw: |
| raise ValueError("Unexpected TXT end") |
| text_digest.update(raw) |
| text = raw.decode("utf-8").removesuffix("\n").removesuffix("\r") |
| if not text: |
| raise ValueError("Empty sentence") |
| texts.append(text) |
| positions.append(position) |
| encoded = processor.encode(texts, out_type=int, num_threads=1) |
| decoded = processor.decode(encoded, num_threads=1) |
| source_bytes = source_stream.read(len(texts)) |
| if len(source_bytes) != len(texts): |
| raise ValueError("Source index is shorter than corpus") |
| source_digest.update(source_bytes) |
| batch_tokens, batch_offsets = [], [] |
| for text, ids, restored, source in zip(texts, encoded, decoded, source_bytes, strict=True): |
| if restored != text or any(i in reserved or i >= 8192 for i in ids): |
| raise ValueError(f"Roundtrip/reserved/ID error: {job['name']} row {count}") |
| if source not in reverse_sources: |
| raise ValueError("Unknown source ID") |
| count += 1 |
| characters += len(text) |
| content += len(ids) |
| stored += len(ids) + 2 |
| longest = max(longest, len(ids) + 2) |
| lengths[len(ids)] += 1 |
| byte_tokens += sum(i in byte_ids for i in ids) |
| sources[reverse_sources[source]] += 1 |
| batch_tokens.extend((2, *ids, 3)) |
| batch_offsets.append(stored) |
| write_numbers(token_stream, "H", batch_tokens) |
| write_numbers(offset_stream, "Q", batch_offsets) |
| write_numbers(row_stream, "Q", positions) |
| out_sources.write(source_bytes) |
| if text_stream.tell() != job["end"] or count != job["sentences"]: |
| raise ValueError("TXT partition boundaries/counts changed") |
| sizes = {suffix: path.stat().st_size for suffix, path in files.items() if suffix != "done.json"} |
| if sizes != {"tokens.bin": stored * 2, "offsets.bin": (count + 1) * 8, |
| "rows.bin": count * 8, "sources.bin": count}: |
| raise ValueError("Part file sizes are inconsistent") |
| result = {"name": job["name"], "signature": job["signature"], "split": job["split"], |
| "start": job["start"], "end": job["end"], "first_row": job["first_row"], |
| "sentences": count, "characters": characters, "stored_tokens": stored, |
| "content_tokens": content, "max_sequence_tokens": longest, |
| "primary_sources": dict(sources), "lengths": dict(lengths), |
| "byte_fallback_tokens": byte_tokens, "unknown_tokens": 0, "roundtrip_mismatches": 0, |
| "txt_part_sha256": text_digest.hexdigest(), "source_part_sha256": source_digest.hexdigest(), |
| "files": {suffix: {"bytes": path.stat().st_size, "sha256": file_sha(path)} |
| for suffix, path in files.items() if suffix != "done.json"}} |
| dump_json(files["done.json"], result) |
| return result |
|
|
|
|
| def collect(jobs, workers, monitor): |
| results, pending = {}, [] |
| for job in jobs: |
| cached = cached_part(job) |
| if cached: |
| |
| cached["lengths"] = {int(k): v for k, v in cached["lengths"].items()} |
| results[job["name"]] = cached |
| else: |
| pending.append(job) |
| total = sum(job["sentences"] for job in jobs) |
| done = sum(r["sentences"] for r in results.values()) |
| initial, started = done, time.monotonic() |
| print(f"[encode] cached={len(results)} pending={len(pending)} total_rows={total:,}", flush=True) |
| with ProcessPoolExecutor(max_workers=workers) as pool: |
| futures = {pool.submit(encode_txt_part, job): job for job in pending} |
| for future in as_completed(futures): |
| job = futures[future] |
| result = future.result() |
| results[job["name"]] = result |
| done += result["sentences"] |
| elapsed = time.monotonic() - started |
| speed = (done - initial) / elapsed |
| monitor.update("encoding", completed_sentences=done, total_sentences=total, |
| encoding_sentences_per_second=speed, |
| encoding_eta_seconds=(total - done) / speed) |
| print(f"[encode] {done:,}/{total:,} ({100*done/total:.1f}%) {speed:,.0f} rows/s", flush=True) |
| return results |
|
|
|
|
| def merge_txt_split(output, split, jobs, results, expected, monitor): |
| from vimeml.tokenizer.encode_parallel import merge_split |
| totals, histogram, sources = Counter(), Counter(), Counter() |
| for job in jobs: |
| result = results[job["name"]] |
| for key in ("sentences", "characters", "content_tokens", "stored_tokens", "byte_fallback_tokens"): |
| totals[key] += result[key] |
| histogram.update(result["lengths"]) |
| sources.update(result["primary_sources"]) |
| if (totals["sentences"] != expected["sentences"] or totals["characters"] != expected["characters"] |
| or dict(sources) != expected["primary_sources"]): |
| raise ValueError(f"Full split counts/source counts changed: {split}") |
| |
| |
| exported = merge_split(None, output, split, jobs, results, |
| {**dict(totals), "tokens_without_special_tokens": totals["content_tokens"], |
| "tokens_with_bos_eos_per_sentence": totals["stored_tokens"]}, monitor, 0, |
| totals["sentences"], verification="metadata") |
| from vimeml.tokenizer.train import percentile |
| exported.update(unknown_tokens=0, roundtrip_mismatches=0, byte_fallback_tokens=totals["byte_fallback_tokens"], |
| byte_fallback_token_fraction=totals["byte_fallback_tokens"] / totals["content_tokens"], |
| characters_per_token=totals["characters"] / totals["content_tokens"], |
| mean_tokens_per_sentence=totals["content_tokens"] / totals["sentences"], |
| length_percentiles_without_special_tokens={f"p{n}": percentile(histogram, n / 100) for n in (50, 95, 99)}, |
| fits_context_with_bos_eos={"128": {"sentences": sum(n for length, n in histogram.items() if length + 2 <= 128), |
| "fraction": sum(n for length, n in histogram.items() if length + 2 <= 128) / totals["sentences"]}}) |
| return exported |
|
|
|
|
| def prepare_v4_indexes(token_dir, index_dir, checks): |
| from vimeml.training.data import SentenceWindowDataset, prepare_split |
| from vimeml.training.check_data import check_core, check_torch |
| token_sha = file_sha(token_dir / "manifest.json") |
| if (index_dir / "manifest.json").exists(): |
| existing = read_json(index_dir / "manifest.json") |
| if existing.get("token_manifest_sha256") != token_sha or existing.get("context_length") != 128: |
| raise ValueError("Index directory belongs to another token dataset") |
| index_dir.mkdir(parents=True, exist_ok=True) |
| splits = {} |
| with ThreadPoolExecutor(max_workers=2) as pool: |
| pending = {split: pool.submit(prepare_split, token_dir, index_dir, split, 128, token_sha, "sha256") for split in SPLITS} |
| for split, future in pending.items(): |
| splits[split] = future.result() |
| index = {"status": "complete", "format": "vimeml_sentence_windows_v1", "context_length": 128, |
| "token_manifest_sha256": token_sha, "token_file_verification": "sha256", |
| "available_splits": list(SPLITS), "splits": splits, |
| "policy": "Separate sentences; nonoverlapping target windows with shared boundary input; all prediction pairs retained."} |
| dump_json(index_dir / "manifest.json", index) |
| report = {"status": "passed", "core_checks": {}, "pytorch_checks": {}, "config": checks} |
| for split in SPLITS: |
| report["core_checks"][split] = check_core(SentenceWindowDataset(token_dir, index_dir, split), checks["seed"]) |
| report["pytorch_checks"][split] = check_torch(SentenceWindowDataset(token_dir, index_dir, split), checks) |
| dump_json(index_dir / "loader-check.json", report) |
| return index, report |
|
|
|
|
| def run(config_path, resume=False, dry_run=False): |
| import sentencepiece as spm |
| config_path = permitted_path(config_path) |
| config = tomllib.loads(config_path.read_text(encoding="utf-8")) |
| paths = {key: permitted_path(ROOT / value) for key, value in config["paths"].items()} |
| corpus, tokenizer, source_store = (paths[k] for k in ("corpus_dir", "tokenizer_dir", "source_store_dir")) |
| output, index_dir = paths["output_dir"], paths["index_dir"] |
| encoding, checks = config["encoding"], config["checks"] |
| if config["windows"]["context_length"] != 128: |
| raise ValueError("V4 context is fixed to 128") |
| for key in ("workers", "batch_size", "chunk_mib"): |
| if type(encoding[key]) is not int or encoding[key] < 1: |
| raise ValueError(f"Invalid encoding resource: {key}") |
| cm, cs, integrity = (read_json(corpus / name) for name in ("manifest.json", "stats.json", "integrity-check.json")) |
| tm, ts = (read_json(tokenizer / name) for name in ("manifest.json", "stats.json")) |
| experiment = read_json(paths["tokenizer_experiment_dir"] / "manifest.json") |
| sm = read_json(source_store / "manifest.json") |
| if any(m.get("status") != "complete" for m in (cm, tm, experiment, sm)) or integrity.get("status") != "passed": |
| raise ValueError("Complete frozen input manifests required") |
| model = tokenizer / "tokenizer.model" |
| if spm.__version__ != tm["sentencepiece_version"] or file_sha(model) != tm["model_sha256"]: |
| raise ValueError("Tokenizer hash/version changed") |
| if Path(experiment["primary_model"]).resolve() != model.resolve(): |
| raise ValueError("Use the selected V4 8K model") |
| kana = tm["trainer_options"]["required_chars"] |
| verify_model(model, 8192, kana) |
| corpus_sha = file_sha(corpus / "manifest.json") |
| if corpus_sha != integrity["manifest_sha256"] or corpus_sha != recorded_hash(sm, "manifest.json", corpus.name): |
| raise ValueError("Source provenance belongs to another corpus") |
| expected_hashes = {"train": experiment["sample"]["train_sha256"], "validation": experiment["validation_sha256"]} |
| for split in SPLITS: |
| if (expected_hashes[split] != integrity["exports"][f"{split}.txt"]["sha256"] |
| or expected_hashes[split] != recorded_hash(sm, f"{split}.txt") |
| or any(sm["splits"][split][key] != cs["splits"][split][key] for key in ("sentences", "characters", "primary_sources"))): |
| raise ValueError("Source and tokenizer corpus identity mismatch") |
| plan = {"vocab_size": 8192, "context_length": 128, "available_splits": list(SPLITS), |
| "sentences": {s: cs["splits"][s]["sentences"] for s in SPLITS}, |
| "paths": {k: str(v) for k, v in paths.items()}, "encoding": encoding, |
| "test_and_blind_access": False, "model_training_started": False, "dry_run": dry_run} |
| print(json.dumps(plan, ensure_ascii=False, indent=2), flush=True) |
| if dry_run: |
| return plan |
| if output.exists() and any(output.iterdir()) and not resume: |
| raise FileExistsError("Output is nonempty; choose a new path or --resume") |
| if resume and not output.exists(): |
| raise FileNotFoundError("Cannot resume an absent token store") |
| output.mkdir(parents=True, exist_ok=True) |
| with output_lock(output): |
| if (output / "manifest.json").exists(): |
| complete = read_json(output / "manifest.json") |
| if complete.get("status") == "complete": |
| if complete["tokenizer_model_sha256"] != tm["model_sha256"]: |
| raise ValueError("Completed output belongs to another tokenizer") |
| for name, digest in complete["output_sha256"].items(): |
| if file_sha(output / name) != digest: |
| raise ValueError(f"Completed output checksum mismatch: {name}") |
| return prepare_v4_indexes(output, index_dir, checks) |
| started = time.monotonic() |
| with RunMonitor(output, ROOT / "runs/tokenizer/corpus-v4-8k-encode" / str(time.time_ns()), True) as monitor: |
| monitor.update("checking_inputs") |
| inputs = [config_path, model, tokenizer / "manifest.json", tokenizer / "stats.json", |
| paths["tokenizer_experiment_dir"] / "manifest.json", source_store / "manifest.json", |
| corpus / "manifest.json", corpus / "stats.json", corpus / "integrity-check.json"] |
| inputs += [source_store / f"{split}.sources.bin" for split in SPLITS] |
| inputs += [Path(__file__), Path(__file__).with_name("encode_parallel.py"), Path(__file__).with_name("encode.py"), |
| Path(__file__).with_name("monitor.py"), Path(__file__).with_name("v4.py"), |
| ROOT / "src/vimeml/training/data.py", ROOT / "src/vimeml/training/check_data.py"] |
| inputs += [corpus / f"{split}.txt" for split in SPLITS] |
| metadata = {str(p): {"bytes": p.stat().st_size, "mtime_ns": p.stat().st_mtime_ns} for p in inputs} |
| hashes = {str(p): file_sha(p) for p in inputs if p.suffix != ".txt"} |
| jobs, scans = [], {} |
| for split in SPLITS: |
| source_path = source_store / f"{split}.sources.bin" |
| if (source_path.stat().st_size != cs["splits"][split]["sentences"] |
| or hashes[str(source_path)] != sm["output_sha256"][source_path.name]): |
| raise ValueError("Frozen source index changed") |
| selected, scans[split] = scan_txt(corpus / f"{split}.txt", split, |
| encoding["chunk_mib"] * 1024 * 1024, cs["splits"][split], |
| lambda done, total: monitor.update("checking_inputs", preflight_completed_sentences=done, preflight_total_sentences=total)) |
| hashes[str(corpus / f"{split}.txt")] = scans[split]["sha256"] |
| if scans[split]["sha256"] != expected_hashes[split]: |
| raise ValueError(f"Actual TXT checksum changed: {split}") |
| for i, job in enumerate(selected): |
| job.update(name=f"{split}-{i:05d}", source_input=str(source_path)) |
| jobs.extend(selected) |
| signature = hashlib.sha256(json.dumps({"hashes": hashes, "encoding": encoding, |
| "splits": list(SPLITS)}, sort_keys=True).encode()).hexdigest() |
| parts = output / ".parts" |
| parts.mkdir(exist_ok=True) |
| for job in jobs: |
| job.update(model=str(model), parts=str(parts), batch_size=encoding["batch_size"], |
| signature=signature, source_ids=sm["source_ids"], verification="sha256") |
| dump_json(output / "encoding-plan.json", {"signature": signature, "plan": plan, "scans": scans}) |
| results = collect(jobs, encoding["workers"], monitor) |
| exported = {} |
| for split in SPLITS: |
| selected = [job for job in jobs if job["split"] == split] |
| exported[split] = merge_txt_split(output, split, selected, results, cs["splits"][split], monitor) |
| |
| for key, other in (("content_tokens", "tokens_without_special_tokens"), ("stored_tokens", "tokens_with_bos_eos_per_sentence"), |
| ("byte_fallback_tokens", "byte_fallback_tokens")): |
| if exported["validation"][key] != ts["splits"]["validation"][other]: |
| raise ValueError("Full validation token counts changed since tokenizer acceptance") |
| monitor.update("verifying_inputs") |
| if any(file_sha(p) != hashes[str(p)] for p in inputs): |
| raise ValueError("Inputs changed during encoding") |
| dump_json(output / "stats.json", {"actual_vocab_size": 8192, "splits": exported}) |
| artifacts = [output / f"{s}.{suffix}" for s in SPLITS for suffix in ("tokens.bin", "offsets.bin", "rows.bin", "sources.bin")] |
| artifacts.append(output / "stats.json") |
| output_hashes = {p.name: file_sha(p) for p in artifacts} |
| for split in SPLITS: |
| if output_hashes[f"{split}.sources.bin"] != sm["output_sha256"][f"{split}.sources.bin"]: |
| raise ValueError("Source rows were reordered during encoding") |
| dump_json(output / "part-audit.json", {"signature": signature, "parts": [results[job["name"]] for job in jobs]}) |
| dump_json(output / "manifest.json", {"status": "complete", "format": FORMAT, "vocab_size": 8192, |
| "token_dtype": "uint16_le", "offset_dtype": "uint64_le", "offset_unit": "tokens", |
| "special_ids": tm["special_ids"], "reserved_marker_ids": tm["marker_ids"], |
| "sentencepiece_version": spm.__version__, "tokenizer_model_sha256": tm["model_sha256"], |
| "tokenizer_dir": str(tokenizer), "corpus_dir": str(corpus), "available_splits": list(SPLITS), |
| "splits": exported, "source_ids": sm["source_ids"], "input_sha256": hashes, |
| "output_sha256": output_hashes, "input_metadata": metadata, "signature": signature, |
| "sequence_policy": "BOS + complete sentence + EOS; no truncation, concatenation or task markers.", |
| "row_offset_source": "corpus_txt", "row_offset_dtype": "uint64_le", |
| "provenance_policy": "rows.bin stores uint64 byte offsets into the corresponding frozen TXT, not the removed JSONL. sources.bin preserves verified original primary source IDs and row order.", |
| "input_verification": "Full SHA256 before/after; every encoded text roundtrip; all sequence IDs checked; merged offsets/sizes/counts checked.", |
| "test_and_blind_access": False, "corpus_ready_for_lm_training": cm.get("ready_for_lm_training", False), |
| "corpus_quality_mode": cm.get("quality_mode"), "elapsed_seconds": time.monotonic() - started, |
| "purpose": "V4 fixed 8K base LM data; capacity redistribution experiment, not vocabulary speed optimization."}) |
| |
| if parts.resolve().parent != output.resolve(): |
| raise ValueError("Unexpected temporary part path") |
| for job in jobs: |
| for path in part_paths(job).values(): |
| if path.resolve().parent != parts.resolve(): |
| raise ValueError("Unexpected temporary file path") |
| path.unlink() |
| parts.rmdir() |
| monitor.update("complete") |
| index, report = prepare_v4_indexes(output, index_dir, checks) |
| print(json.dumps({"tokens": exported, "windows": index["splits"], "loader_checks": report}, ensure_ascii=False, indent=2), flush=True) |
| return index, report |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--config", type=Path, default=ROOT / "configs/data-v4.toml") |
| parser.add_argument("--resume", action="store_true") |
| parser.add_argument("--dry-run", action="store_true") |
| args = parser.parse_args() |
| run(args.config, args.resume, args.dry_run) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
