src/hf_export.py

"""Hugging Face dataset export: audio layout, metadata, card, optional push."""

from __future__ import annotations

import csv
import json
import shutil
from pathlib import Path
from typing import Dict, List, Optional

from .builder import BuildResult
from .config import DatasetConfig


def _dataset_card(config: DatasetConfig, result: BuildResult, repo_id: str) -> str:
    labels_table = "\n".join(
        f"| `{label}` | {count} |" for label, count in sorted(result.label_counts.items())
    )
    return f"""---
license: cc-by-4.0
pretty_name: {config.dataset_name} wake word synthetic speech dataset
task_categories:
- audio-classification
tags:
- audio
- speech
- wake-word
- keyword-spotting
- text-to-speech
- synthetic-data
- edge-impulse
- tinyml
size_categories:
- n<1K
---

# {config.dataset_name} — Wake Word Synthetic Speech Dataset

Synthetic, augmented audio for training a small wake-word / keyword-spotting
model. Generated with **{result.backend_source}** and local audio augmentation.

## Classes

| Label | Samples |
|---|---|
{labels_table}

- `{config.wake_label}` — the target wake phrase and close variants.
- `{config.unknown_label}` — near-miss and unrelated short phrases.
- `{config.noise_label}` — synthetic background noise.

## Audio Specification

| Property | Value |
|---|---|
| Format | WAV |
| Channels | Mono |
| Sample rate | {config.sample_rate_hz} Hz |
| Clip length | {config.duration_seconds} seconds |

## Layout

```text
audio/
  train/  {config.wake_label}.<id>.wav ...
  test/   {config.wake_label}.<id>.wav ...
metadata.csv
hf_metadata.csv
selected_voices.csv
dataset_summary.json
```

## Loading

```python
from datasets import load_dataset, Audio

ds = load_dataset("{repo_id}")
ds = ds.cast_column("audio", Audio(sampling_rate={config.sample_rate_hz}))
print(ds)
```

## Edge Impulse

Filenames follow the Edge Impulse label-prefix convention
(`{config.wake_label}.<id>.wav`) so they upload directly:

```bash
edge-impulse-uploader --category training audio/train/*.wav
edge-impulse-uploader --category testing audio/test/*.wav
```

## Limitations

Synthetic TTS is a bootstrap, not a production benchmark. Add real device
and environment recordings before deploying a wake-word product.

## License

CC BY 4.0. Verify that your use of the generated synthetic speech complies
with the terms of the voice models and tools used to create it.
"""


def export_hf_dataset(
    config: DatasetConfig,
    result: BuildResult,
    hf_dir: str,
    repo_id: str = "your-username/your-dataset",
) -> str:
    """Build a Hugging Face-ready folder from a completed build. Returns its path."""
    source_dir = Path(result.out_dir)
    hf_path = Path(hf_dir).resolve()
    if hf_path.exists():
        shutil.rmtree(hf_path)
    (hf_path / "audio" / "train").mkdir(parents=True, exist_ok=True)
    (hf_path / "audio" / "test").mkdir(parents=True, exist_ok=True)

    for src_split, hf_split in (("training", "train"), ("testing", "test")):
        src = source_dir / "edge_impulse_upload" / src_split
        dst = hf_path / "audio" / hf_split
        for wav in sorted(src.glob("*.wav")):
            shutil.copy2(wav, dst / wav.name)

    for name in ("metadata.csv", "selected_voices.csv", "dataset_summary.json"):
        src = source_dir / name
        if src.exists():
            shutil.copy2(src, hf_path / name)

    # HF metadata.csv-style index (audio path + label).
    hf_rows: List[Dict[str, str]] = []
    for split in ("train", "test"):
        for wav in sorted((hf_path / "audio" / split).glob("*.wav")):
            hf_rows.append(
                {
                    "audio": str(wav.relative_to(hf_path)),
                    "label": wav.name.split(".")[0],
                    "split": split,
                    "filename": wav.name,
                }
            )
    with (hf_path / "hf_metadata.csv").open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["audio", "label", "split", "filename"])
        writer.writeheader()
        writer.writerows(hf_rows)

    (hf_path / "README.md").write_text(
        _dataset_card(config, result, repo_id), encoding="utf-8"
    )
    (hf_path / "hf_dataset_summary.json").write_text(
        json.dumps({"total_wavs": len(hf_rows), "repo_id": repo_id}, indent=2),
        encoding="utf-8",
    )
    return str(hf_path)


def push_to_hub(hf_dir: str, repo_id: str, token: str, private: bool = False) -> str:
    """Upload the HF dataset folder to the Hub. Returns the dataset URL."""
    from huggingface_hub import HfApi

    api = HfApi(token=token)
    api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True, private=private)
    api.upload_folder(
        folder_path=hf_dir,
        repo_id=repo_id,
        repo_type="dataset",
    )
    return f"https://huggingface.co/datasets/{repo_id}"