app.py

"""Gradio Hugging Face Space: wake-word dataset creator.

Generates a keyword-spotting dataset using Google Cloud TTS when an API
key is supplied, and automatically falls back to free Piper TTS otherwise.
Optionally pushes the result to a Hugging Face dataset repo and/or uploads
directly to an Edge Impulse project.
"""

from __future__ import annotations

import os
import shutil
import tempfile
from pathlib import Path
from typing import List, Optional

import gradio as gr

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

# API keys can also be provided as Space secrets.
ENV_GCP_KEY = os.environ.get("GCP_TTS_API_KEY", "")
ENV_HF_TOKEN = os.environ.get("HF_TOKEN", "")
ENV_EI_KEY = os.environ.get("EDGE_IMPULSE_API_KEY", "")


def _split_lines(text: str, fallback: List[str]) -> List[str]:
    items = [line.strip() for line in (text or "").splitlines() if line.strip()]
    return items or list(fallback)


def create_dataset(
    dataset_name: str,
    wake_label: str,
    wake_phrases_text: str,
    unknown_phrases_text: str,
    gcp_api_key: str,
    base_repeats: int,
    augmentations: int,
    background_noise: int,
    max_voices: int,
    test_ratio: float,
    hf_repo_id: str,
    hf_token: str,
    hf_private: bool,
    do_push_hf: bool,
    ei_api_key: str,
    do_upload_ei: bool,
    ei_allow_duplicates: bool,
    progress=gr.Progress(track_tqdm=False),
):
    logs: List[str] = []

    def log(message: str) -> str:
        logs.append(message)
        return "\n".join(logs)

    work_root = Path(tempfile.mkdtemp(prefix="wakeword_"))
    dataset_dir = work_root / "dataset"
    hf_dir = work_root / "hf_dataset"

    gcp_api_key = (gcp_api_key or "").strip() or ENV_GCP_KEY

    try:
        progress(0.05, desc="Selecting TTS backend")
        log("Selecting TTS backend...")
        backend = select_backend(
            gcp_api_key=gcp_api_key,
            language_prefixes=["en", "nl", "de", "fr", "es"],
            max_gcp_voices_per_locale=3,
            max_piper_voices=int(max_voices),
            sample_rate_hz=16000,
        )
        engine = (
            "Google Cloud TTS"
            if backend.source == "google_cloud_tts"
            else "Piper TTS (free fallback)"
        )
        yield log(f"Using backend: {engine}"), None, None

        config = DatasetConfig(
            out_dir=str(dataset_dir),
            dataset_name=dataset_name or "hey_android",
            wake_label=wake_label or "hey_android",
            wake_phrases=_split_lines(wake_phrases_text, DEFAULT_WAKE_PHRASES),
            unknown_phrases=_split_lines(unknown_phrases_text, DEFAULT_UNKNOWN_PHRASES),
            base_repeats_per_phrase_per_voice=int(base_repeats),
            augmentations_per_speech_clip=int(augmentations),
            background_noise_samples=int(background_noise),
            max_piper_voices=int(max_voices),
            test_ratio=float(test_ratio),
        )

        progress(0.15, desc="Generating audio")
        result = build_dataset(config, backend, progress=lambda m: logs.append(m))
        yield log(
            f"Generated {result.total_samples} samples "
            f"(base={result.generated_base}, augmented={result.generated_augmented}, "
            f"failed={result.failed})."
        ), None, None

        progress(0.7, desc="Preparing Hugging Face layout")
        export_hf_dataset(config, result, str(hf_dir), repo_id=hf_repo_id or "your-username/your-dataset")
        log("Hugging Face dataset folder prepared.")

        # Zip for download.
        zip_base = work_root / f"{config.dataset_name}_hf_dataset"
        zip_path = shutil.make_archive(str(zip_base), "zip", str(hf_dir))
        yield log(f"Created download archive: {Path(zip_path).name}"), zip_path, None

        # Optional: push to Hugging Face Hub.
        token = (hf_token or "").strip() or ENV_HF_TOKEN
        if do_push_hf:
            if not token:
                log("Skipping HF push: no token provided.")
            elif not hf_repo_id or "/" not in hf_repo_id:
                log("Skipping HF push: provide a repo id like 'username/dataset-name'.")
            else:
                progress(0.85, desc="Pushing to Hugging Face")
                log(f"Pushing to Hugging Face dataset '{hf_repo_id}'...")
                url = push_to_hub(str(hf_dir), hf_repo_id, token, private=bool(hf_private))
                log(f"Pushed: {url}")
                yield "\n".join(logs), zip_path, None

        # Optional: upload to Edge Impulse.
        ei_key = (ei_api_key or "").strip() or ENV_EI_KEY
        if do_upload_ei:
            if not ei_key:
                log("Skipping Edge Impulse upload: no API key provided.")
            else:
                progress(0.92, desc="Uploading to Edge Impulse")
                log("Uploading dataset to your Edge Impulse project...")
                ei_result = edge_impulse.upload_dataset(
                    dataset_dir=str(dataset_dir),
                    api_key=ei_key,
                    allow_duplicates=bool(ei_allow_duplicates),
                    progress=lambda m: logs.append(m),
                )
                log(
                    f"Edge Impulse: {ei_result.uploaded} uploaded, {ei_result.failed} failed."
                )
                if ei_result.errors:
                    log("Edge Impulse errors:\n" + "\n".join(ei_result.errors[:5]))

        progress(1.0, desc="Done")
        summary = (
            f"### Done\n"
            f"- Backend: **{engine}**\n"
            f"- Total samples: **{result.total_samples}**\n"
            + "\n".join(f"- `{k}`: {v}" for k, v in sorted(result.label_counts.items()))
        )
        yield "\n".join(logs), zip_path, summary

    except Exception as exc:  # noqa: BLE001 - surface errors to the UI
        log(f"ERROR: {exc}")
        yield "\n".join(logs), None, f"### Failed\n\n```\n{exc}\n```"


with gr.Blocks(title="WakeForge — GCP & Piper TTS Wake Word Dataset Creator") as demo:
    gr.Markdown(
        """
        # 🔨 WakeForge
        ### GCP & Piper TTS Wake Word Dataset Creator
        Generate a keyword-spotting dataset for **Hugging Face** and **Edge Impulse**.

        - Provide a **Google Cloud TTS API key** to use Google voices.
        - **No key? It automatically falls back to free Piper TTS.**
        - Optionally **push to a Hugging Face dataset** and/or **upload to your Edge Impulse project**.
        """
    )

    with gr.Row():
        with gr.Column():
            gr.Markdown("### Dataset")
            dataset_name = gr.Textbox(label="Dataset name", value="hey_android")
            wake_label = gr.Textbox(label="Wake label", value="hey_android")
            wake_phrases_text = gr.Textbox(
                label="Wake phrases (one per line)",
                value="\n".join(DEFAULT_WAKE_PHRASES),
                lines=6,
            )
            unknown_phrases_text = gr.Textbox(
                label="Unknown / near-miss phrases (one per line)",
                value="\n".join(DEFAULT_UNKNOWN_PHRASES),
                lines=8,
            )

            gr.Markdown("### Size")
            base_repeats = gr.Slider(1, 5, value=1, step=1, label="Base clips per phrase per voice")
            augmentations = gr.Slider(0, 20, value=8, step=1, label="Augmentations per clip")
            background_noise = gr.Slider(0, 500, value=200, step=10, label="Background noise clips")
            max_voices = gr.Slider(1, 7, value=7, step=1, label="Max voices")
            test_ratio = gr.Slider(0.05, 0.5, value=0.2, step=0.05, label="Test split ratio")

        with gr.Column():
            gr.Markdown("### Google Cloud TTS (optional)")
            gcp_api_key = gr.Textbox(
                label="GCP TTS API key",
                type="password",
                placeholder="Leave blank to use free Piper TTS",
            )

            gr.Markdown("### Push to Hugging Face (optional)")
            do_push_hf = gr.Checkbox(label="Push dataset to Hugging Face Hub", value=False)
            hf_repo_id = gr.Textbox(label="HF dataset repo id", placeholder="username/dataset-name")
            hf_token = gr.Textbox(label="HF write token", type="password", placeholder="hf_...")
            hf_private = gr.Checkbox(label="Private dataset", value=False)

            gr.Markdown("### Upload to Edge Impulse (optional)")
            do_upload_ei = gr.Checkbox(label="Upload dataset to Edge Impulse project", value=False)
            ei_api_key = gr.Textbox(
                label="Edge Impulse API key",
                type="password",
                placeholder="ei_... (Project → Dashboard → Keys)",
            )
            ei_allow_duplicates = gr.Checkbox(label="Allow duplicate samples", value=False)

    generate_btn = gr.Button("Generate dataset", variant="primary")

    summary_md = gr.Markdown()
    download = gr.File(label="Download dataset (zip)")
    logs_box = gr.Textbox(label="Logs", lines=16, max_lines=30)

    generate_btn.click(
        fn=create_dataset,
        inputs=[
            dataset_name, wake_label, wake_phrases_text, unknown_phrases_text,
            gcp_api_key, base_repeats, augmentations, background_noise, max_voices, test_ratio,
            hf_repo_id, hf_token, hf_private, do_push_hf,
            ei_api_key, do_upload_ei, ei_allow_duplicates,
        ],
        outputs=[logs_box, download, summary_md],
    )


if __name__ == "__main__":
    demo.queue().launch()