Ship the data scripts with the checkpoint

tts_data.py

@@ -0,0 +1,267 @@
+"""Turn one PLD language into SpeechT5 training inputs, cached on the Hub.
+
+Why this exists: `halohalo/finetune_tts.py` reads PLD from a raw corpus on a
+local disk (`PLD_RAW`), which a Colab VM does not have. The Hub copy
+(`sapinsapin/pld`) is 201 train shards of every language interleaved, so
+pulling one language means touching all ~24 GB.
+
+That is a one-time cost, and this script pays it once: stream the shards one at
+a time (never more than one on disk), keep the rows for one language, run the
+same text/mel/x-vector preprocessing the original training used, and write the
+result to parquet. Push that to a Hub dataset repo and every later run --
+including every resume after a free-tier VM disappears -- pulls ~1 GB instead
+of 24.
+
+Preprocessing is deliberately identical to `finetune_tts.py`:
+  - text: curly apostrophes normalized, digit-bearing lines dropped
+  - audio: SpeechT5Processor log-mel targets at 16 kHz
+  - speaker: one speechbrain x-vector per *clip*, F.normalize'd, never averaged
+  - caps: <=220 input ids, <=960 mel frames
+
+Usage (on the VM):
+    python scripts/tts_data.py --language ceb --push-to Splintir/pld-ceb-tts-proc
+"""
+
+from __future__ import annotations
+
+import argparse
+import io
+import os
+import re
+import shutil
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+
+SR = 16000
+PLD_REPO = "sapinsapin/pld"
+N_TRAIN_SHARDS = 201
+DIGIT_RE = re.compile(r"\d")
+
+# Columns read from each shard. `audio` dominates the transfer; the rest are
+# cheap and decide whether a row is kept at all.
+COLUMNS = [
+    "audio",
+    "sentence",
+    "language",
+    "speech_type",
+    "num_words",
+    "text_is_prompt",
+    "speaker_id",
+    "duration",
+]
+
+MAX_INPUT_IDS = 220
+MAX_MEL_FRAMES = 960
+
+
+def clean_text(text: str) -> str | None:
+    """SpeechT5's tokenizer is character-level Latin: normalize apostrophes and
+    reject digits rather than teach the model to skip them."""
+    text = text.replace("’", "'").replace("‘", "'").strip()
+    if not text or DIGIT_RE.search(text):
+        return None
+    return text
+
+
+def keep_row(row: dict, language: str) -> bool:
+    """The TTS filter from halolib.finetune._PLD_FILTERS, plus the language."""
+    return (
+        row["language"] == language
+        and row["speech_type"] == "read"
+        and not row["text_is_prompt"]
+        and row["num_words"] >= 3
+    )
+
+
+def build_embedder():
+    import torch
+
+    # speechbrain 1.1 registers optional integrations (k2, wordemb, ...) as lazy
+    # modules that import on any attribute access. Loading the Xvector lobe goes
+    # through pydoc, which probes dunders on every module in sys.modules and so
+    # force-imports integrations whose dependencies are absent. Dunders never
+    # come from a lazy import, so refusing them is safe. Same patch as
+    # sapin-hil/tts.py.
+    from speechbrain.utils import importutils as _importutils
+
+    _lazy_getattr = _importutils.LazyModule.__getattr__
+    _importutils.LazyModule.__getattr__ = (
+        lambda self, attr: (_ for _ in ()).throw(AttributeError(attr))
+        if attr.startswith("__") else _lazy_getattr(self, attr))
+
+    from speechbrain.inference.speaker import EncoderClassifier
+    from speechbrain.utils.fetching import LocalStrategy
+
+    savedir = Path(os.environ.get("HF_HOME", "~/.cache")).expanduser() / "speechbrain-xvect"
+    return EncoderClassifier.from_hparams(
+        source="speechbrain/spkrec-xvect-voxceleb",
+        savedir=str(savedir),
+        # COPY, not the default SYMLINK: symlinking needs Developer Mode on
+        # Windows and fails with WinError 1314 otherwise. Harmless on Linux.
+        local_strategy=LocalStrategy.COPY,
+        run_opts={"device": "cuda" if torch.cuda.is_available() else "cpu"},
+    )
+
+
+def iter_shard_rows(language: str, token: str | None, shards: range):
+    """Yield matching rows shard by shard, holding one shard on disk at a time."""
+    import pyarrow.parquet as pq
+    from huggingface_hub import hf_hub_download
+
+    scratch = Path("/content/_pld_shard") if Path("/content").exists() else Path("./_pld_shard")
+
+    for i in shards:
+        name = f"data/train-{i:05d}-of-{N_TRAIN_SHARDS:05d}.parquet"
+        shutil.rmtree(scratch, ignore_errors=True)
+        scratch.mkdir(parents=True, exist_ok=True)
+        path = hf_hub_download(
+            PLD_REPO, name, repo_type="dataset", token=token, local_dir=str(scratch)
+        )
+        table = pq.read_table(path, columns=COLUMNS)
+        # to_pylist on the whole shard materializes 1500 audio blobs (~120 MB);
+        # row-group at a time keeps the peak an order of magnitude lower.
+        for batch in table.to_batches(max_chunksize=100):
+            for row in batch.to_pylist():
+                if keep_row(row, language):
+                    yield i, row
+        del table
+        shutil.rmtree(scratch, ignore_errors=True)
+
+
+def process(rows, processor, embedder, log_every: int = 250):
+    """(audio, text) -> (input_ids, mel labels, x-vector), dropping what cannot train."""
+    import soundfile as sf
+    import torch
+
+    kept = seen = 0
+    t0 = time.time()
+    for shard_i, row in rows:
+        seen += 1
+        text = clean_text(row["sentence"])
+        if text is None:
+            continue
+
+        wav, sr = sf.read(io.BytesIO(row["audio"]["bytes"]), dtype="float32")
+        if sr != SR:
+            raise SystemExit(f"expected {SR} Hz, shard {shard_i} gave {sr}")
+        if wav.ndim > 1:
+            wav = wav.mean(axis=1)
+
+        example = processor(
+            text=text, audio_target=wav, sampling_rate=SR, return_attention_mask=False
+        )
+        input_ids = example["input_ids"]
+        labels = np.asarray(example["labels"][0], dtype=np.float32)
+        if len(input_ids) > MAX_INPUT_IDS or len(labels) > MAX_MEL_FRAMES:
+            continue
+
+        with torch.no_grad():
+            emb = embedder.encode_batch(torch.tensor(wav).unsqueeze(0))
+            emb = torch.nn.functional.normalize(emb, dim=2).squeeze().cpu().numpy()
+
+        kept += 1
+        if kept % log_every == 0:
+            rate = seen / max(time.time() - t0, 1e-9)
+            print(
+                f"  shard {shard_i:3d}  scanned {seen:6d}  kept {kept:6d}"
+                f"  ({rate:.1f} rows/s)",
+                flush=True,
+            )
+
+        yield {
+            "input_ids": list(map(int, input_ids)),
+            "labels": labels.reshape(-1).tolist(),
+            "n_mel_frames": int(labels.shape[0]),
+            "speaker_embeddings": emb.astype(np.float32).tolist(),
+            "text": text,
+            "speaker_id": row["speaker_id"],
+            "duration": float(row["duration"]),
+        }
+
+
+def write_parquet(records, out_dir: Path, rows_per_file: int = 1000) -> int:
+    """Stream records to parquet shards so peak memory stays near one shard."""
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+
+    out_dir.mkdir(parents=True, exist_ok=True)
+    buf, n, part = [], 0, 0
+
+    def flush() -> None:
+        nonlocal buf, part
+        if not buf:
+            return
+        pq.write_table(pa.Table.from_pylist(buf), out_dir / f"part-{part:04d}.parquet")
+        part += 1
+        buf = []
+
+    for rec in records:
+        buf.append(rec)
+        n += 1
+        if len(buf) >= rows_per_file:
+            flush()
+    flush()
+    return n
+
+
+def main() -> None:
+    sys.stdout.reconfigure(encoding="utf-8")
+    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    ap.add_argument("--language", required=True, help="PLD ISO 639-3 code, e.g. ceb")
+    ap.add_argument("--out", default="/content/pld_proc", help="local parquet dir")
+    ap.add_argument("--push-to", default="", help="Hub dataset repo to upload to")
+    ap.add_argument("--private", action="store_true")
+    ap.add_argument("--shards", type=int, default=N_TRAIN_SHARDS,
+                    help="how many train shards to scan (fewer = smaller sample)")
+    ap.add_argument("--max-samples", type=int, default=0, help="0 = every usable clip")
+    args = ap.parse_args()
+
+    os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+    token = os.environ.get("HF_TOKEN") or None
+
+    from transformers import SpeechT5Processor
+
+    processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
+    embedder = build_embedder()
+
+    out_dir = Path(args.out) / args.language
+    shutil.rmtree(out_dir, ignore_errors=True)
+
+    print(f"scanning {args.shards} PLD train shards for {args.language} ...", flush=True)
+    rows = iter_shard_rows(args.language, token, range(args.shards))
+    records = process(rows, processor, embedder)
+    if args.max_samples:
+        import itertools
+
+        records = itertools.islice(records, args.max_samples)
+
+    t0 = time.time()
+    n = write_parquet(records, out_dir)
+    print(f"wrote {n} rows to {out_dir} in {(time.time() - t0) / 60:.1f} min")
+    if n == 0:
+        raise SystemExit(f"no usable {args.language} rows -- check the filters")
+
+    size = sum(f.stat().st_size for f in out_dir.glob("*.parquet")) / 1024**3
+    print(f"parquet size: {size:.2f} GB")
+
+    if args.push_to:
+        from huggingface_hub import HfApi
+
+        api = HfApi(token=token)
+        api.create_repo(args.push_to, repo_type="dataset", exist_ok=True,
+                        private=args.private)
+        api.upload_folder(
+            folder_path=str(out_dir),
+            repo_id=args.push_to,
+            repo_type="dataset",
+            path_in_repo=f"data/{args.language}",
+            commit_message=f"Preprocessed {args.language} TTS inputs ({n} clips)",
+        )
+        print(f"pushed: https://huggingface.co/datasets/{args.push_to}")
+
+
+if __name__ == "__main__":
+    main()