67 lines
2.6 KiB
Python
67 lines
2.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Validate every shard of a BookTTS teacher dataset."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
|
|
def audit(args: argparse.Namespace) -> None:
|
|
metadata_path = args.dataset / "dataset.json"
|
|
metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
|
|
token_map = np.load(args.dataset / "token_id_map.npy")
|
|
samples = 0
|
|
tokens = 0
|
|
bytes_total = metadata_path.stat().st_size + (args.dataset / "token_id_map.npy").stat().st_size
|
|
digest = hashlib.sha256()
|
|
for record in metadata.get("shards", []):
|
|
path = args.dataset / record["file"]
|
|
digest.update(path.read_bytes())
|
|
bytes_total += path.stat().st_size
|
|
with np.load(path) as shard:
|
|
ids = shard["original_input_ids"]
|
|
mask = shard["attention_mask"]
|
|
teacher = shard["teacher"]
|
|
if ids.shape != mask.shape or teacher.shape[:2] != ids.shape or teacher.shape[2:] != (768,):
|
|
raise ValueError(f"Incompatible shapes in {path.name}: {ids.shape}, {mask.shape}, {teacher.shape}")
|
|
if ids.shape[0] != record["samples"]:
|
|
raise ValueError(f"Sample count mismatch in {path.name}")
|
|
if ids.size and (ids.min() < 0 or ids.max() >= token_map.shape[0]):
|
|
raise ValueError(f"Token id outside token map in {path.name}")
|
|
if not np.isfinite(teacher).all():
|
|
raise ValueError(f"Non-finite teacher value in {path.name}")
|
|
if not np.all((mask == 0) | (mask == 1)):
|
|
raise ValueError(f"Non-binary attention mask in {path.name}")
|
|
samples += ids.shape[0]
|
|
tokens += int(mask.sum())
|
|
if samples != metadata["samples"]:
|
|
raise ValueError(f"Dataset sample count is {samples}, metadata says {metadata['samples']}")
|
|
report = {
|
|
"dataset": str(args.dataset),
|
|
"samples": samples,
|
|
"shards": len(metadata.get("shards", [])),
|
|
"tokens": tokens,
|
|
"vocab_size": metadata["vocab_size"],
|
|
"token_map_entries": int(token_map.shape[0]),
|
|
"bytes": bytes_total,
|
|
"shards_sha256": digest.hexdigest().upper(),
|
|
"teacher_dtype": metadata.get("teacher_dtype"),
|
|
}
|
|
print(json.dumps(report, ensure_ascii=False, indent=2))
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--dataset", type=Path, required=True)
|
|
return parser.parse_args()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(audit(parse_args()) or 0)
|