Ship the data scripts with the checkpoint

vits_data.py

@@ -0,0 +1,250 @@
+"""Turn one PLD speaker into a VITS training set, cached on the Hub.
+
+    python scripts/vits_data.py --language ceb --survey
+    python scripts/vits_data.py --language ceb --speaker-id top --push-to Splintir/pld-ceb-vits
+
+Why this exists separately from `tts_data.py`: that script computes log-mel
+spectrograms and x-vectors, because SpeechT5 consumes both. VITS consumes raw
+waveform and text and nothing else -- it learns its own alignment and carries
+one baked-in voice, so there is no speaker embedding to compute.
+
+**Single speaker, deliberately.** MMS/VITS checkpoints hold exactly one voice.
+Finetuning a one-voice model on PLD's many speakers averages them into mush,
+which is the most likely reason the single-speaker SpeechT5 `-solo` run beat the
+full-corpus `-v2` run on every statistic. Applying that lesson before the run
+this time rather than after it.
+
+`--survey` prints the speaker distribution and exits, so the choice of speaker
+is made against clip counts and total duration rather than assumed. A VITS
+finetune wants tens of minutes at minimum; if the top speaker is thin, the
+survey says so before any GPU time is spent.
+
+The scan reuses `tts_data.py`'s shard iterator: PLD's train split is 201 shards
+of every language interleaved (~24 GB), so one language means touching all of
+them, one shard on disk at a time.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import sys
+from collections import Counter, defaultdict
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+from tts_data import PLD_REPO, iter_shard_rows  # noqa: E402
+
+DIGIT_RE = re.compile(r"\d")
+
+# MMS tokenizers are character-level and lowercase, with a per-language vocab
+# that excludes digits. Rejecting a row is honest; keeping it would teach the
+# model that "1990" is silence.
+KEEP_RE = re.compile(r"[^a-zñáéíóú' -]")
+
+
+def clean_text(text: str) -> str | None:
+    text = text.replace("’", "'").replace("‘", "'")
+    text = text.lower().strip()
+    if not text or DIGIT_RE.search(text):
+        return None
+    text = KEEP_RE.sub("", text)
+    text = re.sub(r"\s+", " ", text).strip()
+    return text or None
+
+
+def survey(language: str, token: str | None, shards: int) -> None:
+    """Print who speaks this language and for how long, then stop."""
+    clips: Counter[str] = Counter()
+    secs: defaultdict[str, float] = defaultdict(float)
+
+    # gender/age are not in tts_data.COLUMNS and the iterator only reads those,
+    # so the survey reports clips and duration -- the two numbers that decide
+    # whether a speaker can carry a finetune.
+    for n, (_, row) in enumerate(iter_shard_rows(language, token, range(shards)), 1):
+        sid = row["speaker_id"]
+        clips[sid] += 1
+        secs[sid] += float(row.get("duration") or 0.0)
+        if n % 2000 == 0:
+            print(f"  ... {n} rows, {len(clips)} speakers", flush=True)
+
+    total = sum(clips.values())
+    print(f"\n{language}: {total} usable clips, {len(clips)} speakers, "
+          f"{sum(secs.values()) / 3600:.1f} h total\n")
+    print(f"{'speaker_id':<26}{'clips':>8}{'minutes':>10}")
+    for sid, n in clips.most_common(20):
+        print(f"{sid:<26}{n:>8}{secs[sid] / 60:>10.1f}")
+    print("\nA VITS finetune wants >= ~30 min from one speaker. Pick from the "
+          "top rows and rerun with --speaker-id.")
+
+
+def cache_dir(language: str, speaker: str) -> Path:
+    return Path("vits_cache") / language / speaker
+
+
+def load_cached(language: str, speaker: str):
+    """Reuse a completed scan. Scanning 201 shards to find ~15 minutes of audio
+    costs ~24 GB of transfer, so it must never be repeated because a later step
+    failed."""
+    manifest = cache_dir(language, speaker) / "manifest.jsonl"
+    if not manifest.exists():
+        return None
+    rows = [json.loads(line) for line in
+            manifest.read_text(encoding="utf-8").splitlines() if line.strip()]
+    rows = [r for r in rows if (cache_dir(language, speaker) / r["file"]).exists()]
+    if not rows:
+        return None
+    print(f"reusing {len(rows)} cached clips from "
+          f"{cache_dir(language, speaker)} (delete it to force a rescan)",
+          flush=True)
+    return rows
+
+
+def collect(language: str, speaker: str, token: str | None, shards: int,
+            max_seconds: float):
+    """Gather one speaker's clips. Returns (records, resolved_speaker_id)."""
+    import io
+
+    import soundfile as sf
+
+    # `top` cannot be resolved until the corpus has been scanned once, so the
+    # first pass counts and the second keeps. Two passes over 24 GB is slow;
+    # buffering every language's audio in RAM instead is worse.
+    if speaker == "top":
+        counts: Counter[str] = Counter()
+        for _, row in iter_shard_rows(language, token, range(shards)):
+            counts[row["speaker_id"]] += 1
+        if not counts:
+            raise SystemExit(f"no usable {language} rows -- check the filters")
+        speaker, n = counts.most_common(1)[0]
+        print(f"resolved `top` -> {speaker} ({n} clips)", flush=True)
+
+    out = cache_dir(language, speaker)
+    out.mkdir(parents=True, exist_ok=True)
+    manifest = (out / "manifest.jsonl").open("w", encoding="utf-8")
+
+    records, total = [], 0.0
+    seen_shard = -1
+    for shard, row in iter_shard_rows(language, token, range(shards)):
+        # One speaker is a handful of clips scattered over 201 shards, so a
+        # clip-count progress line can stay silent for hours. Report the scan
+        # itself instead -- otherwise a live run is indistinguishable from a
+        # hung one.
+        if shard != seen_shard:
+            seen_shard = shard
+            print(f"  shard {shard + 1}/{shards}  kept {len(records)} clips, "
+                  f"{total / 60:.1f} min", flush=True)
+        if row["speaker_id"] != speaker:
+            continue
+        text = clean_text(row["sentence"])
+        if not text:
+            continue
+        raw = row["audio"]["bytes"]
+        # Decode once here rather than trusting the shard's declared duration:
+        # the trainer segments on real sample counts, and a mismatch shows up as
+        # a silent crash deep in the collator.
+        try:
+            wav, rate = sf.read(io.BytesIO(raw), dtype="float32", always_2d=False)
+        except Exception as exc:                      # noqa: BLE001
+            print(f"  skipped unreadable clip: {exc}", flush=True)
+            continue
+        if wav.ndim > 1:
+            wav = wav.mean(axis=1)
+        secs = len(wav) / rate
+        if not 1.0 <= secs <= 15.0:
+            continue
+
+        # Write PLD's own encoded bytes straight through rather than re-encoding
+        # the decoded array: no quality loss, and `datasets` can build an Audio
+        # column from encoded bytes without torchcodec, which it needs for raw
+        # arrays and file paths alike.
+        ext = Path(row["audio"].get("path") or "clip.wav").suffix or ".wav"
+        name = f"{len(records):04d}{ext}"
+        (out / name).write_bytes(raw)
+        manifest.write(json.dumps({"file": name, "text": text,
+                                   "seconds": round(secs, 3)}) + "\n")
+        manifest.flush()
+
+        records.append({"file": name, "text": text, "seconds": round(secs, 3)})
+        total += secs
+        if max_seconds and total >= max_seconds:
+            break
+
+    manifest.close()
+    print(f"\n{speaker}: {len(records)} clips, {total / 60:.1f} min "
+          f"-> cached in {out}", flush=True)
+    if total < 900:
+        print("WARNING: under 15 minutes. Expect a weak finetune -- consider "
+              "pooling a second speaker of the same gender and dialect.",
+              flush=True)
+    return records, speaker
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    ap.add_argument("--language", default="ceb", help="PLD ISO 639-3 code")
+    ap.add_argument("--speaker-id", default="top",
+                    help="`top` picks the speaker with the most clips")
+    ap.add_argument("--survey", action="store_true",
+                    help="print the speaker distribution and exit")
+    ap.add_argument("--shards", type=int, default=201)
+    ap.add_argument("--max-seconds", type=float, default=0,
+                    help="stop after this much audio (0 = no cap)")
+    ap.add_argument("--push-to", default="",
+                    help="Hub dataset repo, e.g. Splintir/pld-ceb-vits")
+    args = ap.parse_args()
+
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        env = Path(__file__).resolve().parent.parent / ".env"
+        if env.exists():
+            for line in env.read_text(encoding="utf-8").splitlines():
+                key, _, value = line.strip().partition("=")
+                if key == "HF_TOKEN" and value:
+                    token = value.strip()
+
+    print(f"scanning {args.shards} {PLD_REPO} train shards for {args.language} ...",
+          flush=True)
+
+    if args.survey:
+        survey(args.language, token, args.shards)
+        return
+
+    speaker = args.speaker_id
+    records = None if speaker == "top" else load_cached(args.language, speaker)
+    if records is None:
+        records, speaker = collect(args.language, args.speaker_id, token,
+                                   args.shards, args.max_seconds)
+    if not records:
+        raise SystemExit(f"no clips for speaker {speaker}")
+
+    from datasets import Audio, Dataset
+
+    src = cache_dir(args.language, speaker)
+    rows = [{"audio": {"bytes": (src / r["file"]).read_bytes(),
+                       "path": r["file"]},
+             "text": r["text"]}
+            for r in records]
+    ds = Dataset.from_list(rows).cast_column("audio", Audio(sampling_rate=16000))
+    # A held-out slice the trainer can score against, kept small: VITS eval is
+    # slow (it renders audio) and the number that decides anything is the
+    # 50-line bench, not this. Proportional with a floor and a ceiling -- a
+    # flat floor alone puts more clips in eval than train on a small set.
+    n_eval = max(4, min(16, len(ds) // 10))
+    ds = ds.train_test_split(test_size=n_eval, seed=0)
+    print(ds)
+
+    if args.push_to:
+        ds.push_to_hub(args.push_to, token=token, private=True)
+        print(f"pushed -> {args.push_to}  (speaker {speaker})")
+    else:
+        out = Path("vits_data") / args.language
+        ds.save_to_disk(str(out))
+        print(f"saved -> {out}  (pass --push-to to upload)")
+
+
+if __name__ == "__main__":
+    main()