build_samples.py

#!/usr/bin/env python3
"""
Build the Audio Regenerate sample gallery for the project page.

Reads the website-sample output layout produced by scripts/make_uvm_test_set.py:

    <src>/output/<name>_chunk<N>/speech_edit.wav
    <src>/output/<name>_chunk<N>/transcript.txt

The transcript marks the edited span with separators:

    original: prefix | original middle | suffix
    edited:   prefix | edited middle   | suffix

Edit SELECTED_SAMPLES below to choose the exact samples, in display order.
For example, ("alejandro", 3) means:

    <src>/output/alejandro_chunk3/

For each selected sample, this script copies both speech_edit.wav and
transcript.txt into assets/samples/ and writes samples.js with the
edited transcript split into display segments.

Usage:
    python3 build_samples.py --src /path/to/samples_website/output

Re-run any time you regenerate samples.
"""
import argparse
import json
import re
import shutil
from pathlib import Path

HERE = Path(__file__).resolve().parent

# Ordered list of samples to include on the website.
#
# Each tuple is (<file name / speaker prefix>, <chunk number>), matching folders
# like samples_website/output/output/alejandro_chunk3.
SELECTED_SAMPLES = [
    ("stephen", 0),
    ("sumukh", 0),
    ("xingzhe", 4),
    ("alejandro", 3),
    ("mithilesh", 0),
    ("jasmine", 2),
]


def norm(s: str) -> str:
    return re.sub(r"\s+", " ", (s or "").strip())


def segments(prefix: str, middle: str, suffix: str):
    """Build display segments; the middle is the regenerated/edited span."""
    out = []
    if norm(prefix):
        out.append({"text": norm(prefix), "edited": False})
    if norm(middle):
        out.append({"text": norm(middle), "edited": True})
    if norm(suffix):
        out.append({"text": norm(suffix), "edited": False})
    return out


def parse_transcript(path: Path):
    """Return edited transcript segments from output/<id>/transcript.txt."""
    lines = path.read_text(encoding="utf-8", errors="ignore").splitlines()
    edited_line = ""
    for line in lines:
        if line.lower().startswith("edited:"):
            edited_line = line.split(":", 1)[1]
            break
    if not edited_line:
        raise ValueError(f"Could not find an 'edited:' line in {path}")

    parts = [norm(p) for p in edited_line.split("|")]
    if len(parts) < 3:
        raise ValueError(f"Expected edited transcript with two '|' separators in {path}")

    prefix = parts[0]
    middle = " ".join(p for p in parts[1:-1] if p)
    suffix = parts[-1]
    return segments(prefix, middle, suffix)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", required=True,
                    help="Impute output dir (e.g. speech-edit/samples_website/output).")
    args = ap.parse_args()

    src = Path(args.src)
    sample_root = src / "output" if (src / "output").is_dir() else src
    if not sample_root.is_dir():
        raise SystemExit(f"Expected {sample_root} to exist.")

    dst_dir = HERE / "assets" / "samples"
    dst_dir.mkdir(parents=True, exist_ok=True)
    for old in dst_dir.glob("sample-*"):
        old.unlink()

    samples = []
    for i, (name, chunk) in enumerate(SELECTED_SAMPLES, start=1):
        source_id = f"{name}_chunk{chunk}"
        source_dir = sample_root / source_id
        edit_wav = source_dir / "speech_edit.wav"
        transcript = source_dir / "transcript.txt"
        missing = [str(p) for p in (edit_wav, transcript) if not p.exists()]
        if missing:
            raise SystemExit(f"Missing files for {source_id}: {', '.join(missing)}")

        edit_dst = dst_dir / f"sample-{i:02d}.wav"
        transcript_dst = dst_dir / f"sample-{i:02d}_transcript.txt"
        shutil.copyfile(edit_wav, edit_dst)
        shutil.copyfile(transcript, transcript_dst)

        samples.append({
            "id": f"sample-{i:02d}",
            "title": f"Sample {i:02d}",
            "sourceId": source_id,
            "audio": f"assets/samples/{edit_dst.name}",
            "transcript": f"assets/samples/{transcript_dst.name}",
            "segments": parse_transcript(transcript),
        })

    header = (
        "/* AUTO-GENERATED by build_samples.py — do not edit by hand.\n"
        "   Edit SELECTED_SAMPLES in build_samples.py, then re-run:\n"
        "   python3 build_samples.py --src <dir>\n"
        "   Source layout: <src>/output/<name>_chunk<N>/{speech_edit.wav,transcript.txt}. */\n\n"
    )
    body = "window.SAMPLES = " + json.dumps(samples, indent=2, ensure_ascii=False) + ";\n"
    (HERE / "samples.js").write_text(header + body, encoding="utf-8")

    print(f"Wrote {len(samples)} samples to samples.js")
    for s in samples:
        new = " ".join(seg["text"] for seg in s["segments"] if seg["edited"])
        print(f"  {s['id']} <- {s['sourceId']}: \u201c{new}\u201d")


if __name__ == "__main__":
    main()