File size: 2,722 Bytes
840c3a4 | 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 | #!/usr/bin/env python3
"""Verify processed datasets against pinned integrity manifests."""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
import sys
import yaml
OUTPUT_NAMES = {"lm1b": "lm1b-bert-128", "owt": "owt-gpt2-1024"}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1])
parser.add_argument("--config", type=Path)
parser.add_argument("--dataset", choices=("all", "lm1b", "owt"), default="all")
return parser.parse_args()
def main() -> int:
args = parse_args()
root = args.root.resolve()
sys.path.insert(0, str(root / "src"))
from dlb.data import build_processing_contract, verify_processed_dataset
from dlb.io import atomic_json_write
config_path = args.config or root / "artifacts" / "data.yaml"
configuration = yaml.safe_load(config_path.read_text(encoding="utf-8"))
hf_home = root / "data" / "raw" / "huggingface"
hub_cache = hf_home / "hub"
os.environ["HF_HOME"] = str(hf_home)
from datasets import load_from_disk
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer
names = ("lm1b", "owt") if args.dataset == "all" else (args.dataset,)
for name in names:
contract = build_processing_contract(configuration, name)
tokenizer_name = contract["tokenizer_id"]
tokenizer_revision = contract["tokenizer_revision"]
manifest_path = root / "data" / "manifests" / f"{name}.json"
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
snapshot_path = snapshot_download(
repo_id=tokenizer_name,
revision=tokenizer_revision,
cache_dir=hub_cache,
local_files_only=True,
)
tokenizer = AutoTokenizer.from_pretrained(snapshot_path, local_files_only=True)
output_dir = root / "data" / "processed" / OUTPUT_NAMES[name]
dataset = load_from_disk(str(output_dir))
verified = verify_processed_dataset(
dataset=dataset,
manifest=manifest,
output_dir=output_dir,
tokenizer=tokenizer,
expected_contract=contract,
)
atomic_json_write(manifest_path, verified)
print(
f"OK {name} documents={verified['document_counts']} "
f"packed={verified['packed_sequence_counts']} "
f"checked={verified['verification']['checked_sequences']} "
f"files={verified['verification']['files_checked']} verified=true"
)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|