generate.py

"""Command-line dataset generator.

Examples
--------
Free Piper TTS (no key)::

    python generate.py --out output --push-hf-repo user/hey-android

Google Cloud TTS + Edge Impulse upload::

    python generate.py --gcp-api-key "$GCP_TTS_API_KEY" \\
        --edge-impulse-api-key "$EDGE_IMPULSE_API_KEY"
"""

from __future__ import annotations

import argparse
import os

from src import edge_impulse
from src.backends import select_backend
from src.builder import build_dataset
from src.config import DEFAULT_UNKNOWN_PHRASES, DEFAULT_WAKE_PHRASES, DatasetConfig
from src.hf_export import export_hf_dataset, push_to_hub


def _read_lines(path: str | None, fallback: list[str]) -> list[str]:
    if not path:
        return list(fallback)
    lines = [ln.strip() for ln in open(path, encoding="utf-8") if ln.strip()]
    return lines or list(fallback)


def main() -> None:
    p = argparse.ArgumentParser(description="Generate a wake-word dataset (GCP TTS or free Piper).")
    p.add_argument("--out", default="output")
    p.add_argument("--hf-out", default="hf_dataset")
    p.add_argument("--dataset-name", default="hey_android")
    p.add_argument("--wake-label", default="hey_android")
    p.add_argument("--wake-phrases-file", default=None)
    p.add_argument("--unknown-phrases-file", default=None)

    p.add_argument("--gcp-api-key", default=os.environ.get("GCP_TTS_API_KEY", ""))
    p.add_argument("--base-repeats", type=int, default=1)
    p.add_argument("--augmentations", type=int, default=8)
    p.add_argument("--background-noise", type=int, default=200)
    p.add_argument("--max-voices", type=int, default=7)
    p.add_argument("--test-ratio", type=float, default=0.2)

    p.add_argument("--push-hf-repo", default=None, help="e.g. username/dataset-name")
    p.add_argument("--hf-token", default=os.environ.get("HF_TOKEN", ""))
    p.add_argument("--hf-private", action="store_true")

    p.add_argument("--edge-impulse-api-key", default=os.environ.get("EDGE_IMPULSE_API_KEY", ""))
    p.add_argument("--ei-allow-duplicates", action="store_true")

    args = p.parse_args()

    backend = select_backend(
        gcp_api_key=args.gcp_api_key,
        language_prefixes=["en", "nl", "de", "fr", "es"],
        max_gcp_voices_per_locale=3,
        max_piper_voices=args.max_voices,
        sample_rate_hz=16000,
    )
    print(f"Backend: {backend.source}")

    config = DatasetConfig(
        out_dir=args.out,
        dataset_name=args.dataset_name,
        wake_label=args.wake_label,
        wake_phrases=_read_lines(args.wake_phrases_file, DEFAULT_WAKE_PHRASES),
        unknown_phrases=_read_lines(args.unknown_phrases_file, DEFAULT_UNKNOWN_PHRASES),
        base_repeats_per_phrase_per_voice=args.base_repeats,
        augmentations_per_speech_clip=args.augmentations,
        background_noise_samples=args.background_noise,
        max_piper_voices=args.max_voices,
        test_ratio=args.test_ratio,
    )

    result = build_dataset(config, backend)
    repo_id = args.push_hf_repo or "your-username/your-dataset"
    hf_dir = export_hf_dataset(config, result, args.hf_out, repo_id=repo_id)
    print(f"Hugging Face dataset folder: {hf_dir}")

    if args.push_hf_repo and args.hf_token:
        url = push_to_hub(hf_dir, args.push_hf_repo, args.hf_token, private=args.hf_private)
        print(f"Pushed dataset: {url}")

    if args.edge_impulse_api_key:
        ei_result = edge_impulse.upload_dataset(
            dataset_dir=args.out,
            api_key=args.edge_impulse_api_key,
            allow_duplicates=args.ei_allow_duplicates,
        )
        print(f"Edge Impulse: {ei_result.uploaded} uploaded, {ei_result.failed} failed.")


if __name__ == "__main__":
    main()