diff --git a/CLAUDE.md b/CLAUDE.md index bc4f9f2..e3a559c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -34,6 +34,7 @@ noisekit/ noisekit generate --dataset --samples N --preset P1 --preset P2 --output ./out --seed 42 noisekit generate ... --preset noise --noise-dir /path/to/noise_wavs noisekit generate ... --no-nisqa # skip NISQA (no model download, faster) +noisekit generate ... --transcript-column utterance # override transcript column (default: text/sentence/transcription/normalized_text) noisekit score ./audio_dir [--reference-dir ./ref] [--output scores.json] noisekit score ./audio_dir --no-nisqa # skip NISQA for standalone scoring noisekit list-presets [--verbose] diff --git a/README.md b/README.md index a0eb8ac..3382648 100644 --- a/README.md +++ b/README.md @@ -64,6 +64,18 @@ uvx noisekit generate \ `--preset` is repeatable: pass it once per preset. +If your dataset stores transcripts under a different column name (e.g. `utterance`, `raw_text`, `translation`), use `--transcript-column`: + +```bash +uvx noisekit generate \ + --dataset my-org/my-dataset --split test \ + --samples 100 --preset telecom \ + --transcript-column utterance \ + --output ./out +``` + +By default, noisekit tries these columns in order: `text`, `sentence`, `transcription`, `normalized_text`. An error is raised if none are found and `--transcript-column` is not set. + For `noise`, you can supply your own background-noise WAVs with `--noise-dir` (e.g. [MUSAN](https://www.openslr.org/17/), [DEMAND](https://zenodo.org/record/1227121), or [FSD50K](https://zenodo.org/record/4060432)): ```bash diff --git a/noisekit/cli.py b/noisekit/cli.py index d553514..6e4a87b 100644 --- a/noisekit/cli.py +++ b/noisekit/cli.py @@ -46,6 +46,16 @@ def generate( nisqa: Annotated[ bool, typer.Option("--nisqa/--no-nisqa", help="Compute NISQA scores (downloads ~50 MB model on first use)") ] = True, + transcript_column: Annotated[ + str | None, + typer.Option( + "--transcript-column", + help=( + "Dataset column to use as the transcript. " + "Defaults to the first non-empty value among: text, sentence, transcription, normalized_text." + ), + ), + ] = None, ) -> None: """Generate a degraded speech dataset by applying audio presets to a clean source dataset.""" console.print( @@ -69,6 +79,7 @@ def generate( preset_file=preset_file, noise_dir=noise_dir, nisqa=nisqa, + transcript_column=transcript_column, ) diff --git a/noisekit/dataset.py b/noisekit/dataset.py index 5386e9a..2dd4e8b 100644 --- a/noisekit/dataset.py +++ b/noisekit/dataset.py @@ -54,7 +54,13 @@ def extract_language(sample: dict, config: str | None = None) -> str | None: return None -def extract_audio_and_text(sample: dict) -> tuple[np.ndarray, int, str]: +_FALLBACK_TRANSCRIPT_COLS = ("text", "sentence", "transcription", "normalized_text") + + +def extract_audio_and_text( + sample: dict, + transcript_column: str | None = None, +) -> tuple[np.ndarray, int, str]: audio_field = sample["audio"] raw_bytes = audio_field.get("bytes") @@ -73,11 +79,20 @@ def extract_audio_and_text(sample: dict) -> tuple[np.ndarray, int, str]: if array.ndim == 2: array = array.mean(axis=1) # (samples, channels) → (samples,) for mono-only metrics - text = ( - sample.get("text") - or sample.get("sentence") - or sample.get("transcription") - or sample.get("normalized_text") - or "" - ) + if transcript_column is not None: + if transcript_column not in sample: + raise ValueError( + f"Transcript column '{transcript_column}' not found in dataset. " + f"Available columns: {list(sample.keys())}" + ) + text = sample.get(transcript_column) or "" + else: + text = next((sample[c] for c in _FALLBACK_TRANSCRIPT_COLS if sample.get(c)), None) + if text is None: + raise ValueError( + f"No transcript column found. Tried: {list(_FALLBACK_TRANSCRIPT_COLS)}. " + f"Available columns: {list(sample.keys())}. " + "Use --transcript-column to specify the correct column." + ) + return array, int(sr), str(text).strip() diff --git a/noisekit/pipeline.py b/noisekit/pipeline.py index 7167b56..a0155b7 100644 --- a/noisekit/pipeline.py +++ b/noisekit/pipeline.py @@ -35,6 +35,7 @@ def run_generate( preset_file: Path | None, noise_dir: Path | None = None, nisqa: bool = True, + transcript_column: str | None = None, ) -> None: output_dir = Path(output) audio_dir = output_dir / "audio" @@ -55,7 +56,7 @@ def run_generate( _seen_names: set[str] = set() for i, sample in enumerate(track(raw_samples, description="Generating …")): - ref_array, ref_sr, transcript = extract_audio_and_text(sample) + ref_array, ref_sr, transcript = extract_audio_and_text(sample, transcript_column) language = extract_language(sample, config) ref_16k = _resample_to_16k(ref_array, ref_sr) peak = np.abs(ref_16k).max()