AGENTS.md · speechbrain/speechbrain

Sections

1. Writing code
2. Project structure
3. Architecture concepts
4. Recipe conventions
5. Running tests
6. Git workflow
7. Code style
8. Common pitfalls
9. Pretrained models and inference
10. Key dependencies
11. Links

First lines (shown under the repo's Apache-2.0 license)

# AGENTS.md

SpeechBrain is an open-source PyTorch toolkit for conversational AI (speech recognition, speaker verification, speech enhancement, separation, TTS, spoken language understanding, and more) known for its ease of use and flexibility. It is Apache 2.0 licensed.

## Writing code

- **Minimal try/except**: let errors propagate — silent failures hide bugs. Only catch exceptions for intentional fault tolerance (retries, robustness).
- **Targeted comments**: don't explain your work process or reference old code. Use targeted comments sparingly to clarify ambiguous logic.
- **Zen of Python**: remember the Zen of Python when writing code.
```
Beautiful is better than ugly.
Explicit is better than implicit.
Simple is better than complex.
Complex is better than complicated.
Flat is better than nested.
Sparse is better than dense.
Readability counts.
Special cases aren't special enough to break the rules.
Although practicality beats purity.
Errors should never pass silently.
Unless explicitly silenced.
In the face of ambiguity, refuse the temptation to guess.
There should be one-- and preferably only one --obvious way to do it.
Although that way may not be obvious at first unless you're Dutch.
Now is better than never.
Although never is often better than *right* now.
If the implementation is hard to explain, it's a bad idea.
If the implementation is easy to explain, it may be a good idea.
Namespaces are one honking great idea -- let's do more of those!
```

## Project structure

```
speechbrain/          # Core library (importable as `import speechbrain`)
  core.py             # Brain class — the central training/eval orchestrator
  dataio/             # Data loading, batching, samplers, dataset objects
  nnet/               # Neural network building blocks (RNN, CNN, attention, transformers, losses, etc.)
  lobes/              # Higher-level model components (feature extractors, encoders, full models like CRDNN, wav2vec2)
  decoders/           # CTC, seq2seq beam search, transducer decoders
  processing/         # Signal processing (features, augmentation, multi-mic)
  utils/              # Checkpointing, distributed training, metrics, logging, profiling
  inference/          # Pretrained model interfaces (EncoderClassifier, EncoderDecoderASR, etc.)
  integrations/       # Optional heavy-dependency integrations (Transformers, Whisper, etc.)
  lm/                 # Language model utilities
recipes/              # Training scripts organized as recipes/{dataset}/{task}/{model}/
  {dataset}/{task}/
    {model}/          # Recipe implementation for a specific model/configuration
      train.py        # Training script (subclasses Brain)
      hparams/        # HyperPyYAML config files (train.yaml, etc.)
      extra_requirements.txt  # Recipe-specific pip dependencies (if any)
      README.md       # Results, how to run, pretrained model links
templates/            # Minimal working examples to bootstrap new recipes
tests/
  unittests/          # Unit tests for core library
  integration/        # Integration tests (small end-to-end training runs)
docs/                 # Documentation
  tutorials/          # Jupyter notebooks integrated into ReadTheDocs
tools/                # Maintenance scripts (tutorial cell updater, etc.)
```

## Architecture concepts

### The Brain class (`speechbrain.core.Brain`)

Brain is the central abstraction for all training and evaluation. Every recipe subclasses it and overrides the following methods:

- `compute_forward(batch, stage)` — forward pass, returns predictions
- `compute_objectives(predictions, batch, stage)` — computes loss, logs metrics

The `stage` argument is a `Stage` enum: `TRAIN`, `VALID`, or `TEST` which defines the current stage of the training loop. Brain handles the training loop, checkpointing, distributed training (DDP), gradient accumulation, mixed precision, and logging.

```python
class ASR(sb.Brain):
    def compute_forward(self, batch, stage):
        wavs, lens = batch.sig  # (batch, time), (batch,) relative lengths
        feats = self.hparams.compute_features(wavs)
        feats = self.modules.encoder(feats)
        return self.modules.decoder(feats)

    def compute_objectives(self, predictions, batch, stage):
        tokens, token_lens = batch.tokens
        loss = self.hparams.ctc_cost(predictions, tokens, lens, token_lens)
        if stage != sb.Stage.TRAIN:
            self.cer_metric.append(batch.id, predictions, tokens)
        return loss
```

Key lifecycle methods you can override: `on_stage_start`, `on_stage_end`, `on_fit_batch_end`, `fit_batch`, `evaluate_batch`, `init_optimizers`.

### HyperPyYAML (the YAML config system)

SpeechBrain uses HyperPyYAML, an extended YAML syntax maintained by SpeechBrain at https://github.com/speechbrain/HyperPyYAML. This is NOT plain YAML — it is a declarative system that can instantiate Python objects, resolve references, and perform simple arithmetic. Understanding it is essential.

Key tags:
- `!new:module.ClassName` — instantiates a Python object. Indented keys become constructor kwargs.
- `!name:module.ClassName` — returns a callable (like `functools.partial`) without calling it. Used for optimizers, schedulers.
- `!ref <key>` — references another value in the YAML. Supports `<key[subkey]>`, string interpolation (`<folder>/subdir`), and arithmetic (`<foo> * 2 + 1`).
- `!copy <key>` — deep-copies the referenced object (vs `!ref` which is a shallow reference).
- `!apply:module.function` — calls the function immediately during YAML loading.
- `!tuple` — creates a Python tuple.
- `!PLACEHOLDER` — errors at load time, must be overridden via CLI or overrides dict.

Example pattern from a real recipe:
```yaml
seed: 1234
output_folder: !ref results/asr/<seed>
save_folder: !ref <output_folder>/save

model: !new:speechbrain.lobes.models.CRDNN.CRDNN
    output_size: 40
    cnn_blocks: 2
    dnn_blocks: 2

opt_class: !name:torch.optim.Adam
    lr: 0.001

epoch_counter: !new:speechbrain.utils.epoch_loop.EpochCounter
    limit: !ref <number_of_epochs>
```

**Critical: loading YAML executes arbitrary Python code** — `!new:` will import and instantiate anything. Treat YAML files with the same caution as Python code.

Overrides from CLI:
```bash
python train.py hparams/train.yaml --seed 42 --lr 0.0001 --data_folder /path/to/data --num_epochs=100
```

### Tensor convention

All tensors follow **batch-time-channels** ordering:
- `(batch, time)` for raw waveforms
- `(batch, time, channels)` for features
- `(batch, time, channels, ...)` for multi-channel

Lengths are tracked as **relative lengths** (0.0 to 1.0), representing the fraction of the max length in the batch. This avoids passing absolute lengths and simplifies padding/masking. Example: a batch of 3 signals with lengths [16000, 12000, 8000] has relative lengths [1.0, 0.75, 0.5].

### Data pipeline

The pipeline has three layers: **data manifests** (JSON/CSV) → **DynamicItemDataset** → **Dynamic Item Pipelines**.

**Manifests** are JSON or CSV files containing static items (file paths, transcriptions, speaker IDs, durations). JSON format: `{"utt1": {"wav": "path.flac", "wrd": "HELLO", "spk_id": "spk01", "duration": 3.5}, ...}`. Each recipe provides a preparation script that parses raw datasets into this format. Manifests must include a `duration` field for dynamic batching to work.

**DynamicItemDataset** (`speechbrain.dataio.dataset`) loads a manifest and supports on-the-fly transformations via dynamic items. Dependencies between items are resolved automatically as a DAG. Items are evaluated **lazily** — only items in `set_output_keys` (and their dependencies) are computed.

```python
train_data = DynamicItemDataset.from_csv(csv_path=hparams["train_csv"],
    replacements={"data_root": hparams["data_folder"]})  # replacements substitute placeholders in manifest values
```

**Dynamic Item Pipelines** are functions decorated with `@sb.utils.data_pipeline.takes(...)` / `@sb.utils.data_pipeline.provides(...)`.

```python
@sb.utils.data_pipeline.takes("wav")
@sb.utils.data_pipeline.provides("sig")
def audio_pipeline(wav):
    return sb.dataio.dataio.read_audio(wav)

@sb.utils.data_pipeline.takes("wrd")
@sb.utils.data_pipeline.provides("wrd", "tokens_bos", "tokens_eos", "tokens")