Data Readers

Storage and task interfaces for ASR data pipelines

Overview

Omnilingual ASR uses a modular data pipeline architecture separating storage (how data is read) from task (how data is processed). This enables flexible mixing of different storage backends with preprocessing pipelines.

MixtureParquetStorage

Parquet-based storage implementation with partition weighting and multilingual sampling.

Constructor

MixtureParquetStorage(
    path: Path,
    config: MixtureParquetStorageConfig
)

path · Path · required

Path to parquet dataset directory with language/corpus partitions.

config · MixtureParquetStorageConfig · required

Storage configuration including fragment streaming, loading, and weighting parameters.

Configuration: MixtureParquetStorageConfig

fragment_streaming · FragmentStreamingConfig · required

Controls how parquet fragments (row groups) are streamed:

• seed: Random seed for shuffling
• fragment_shuffle_window: Window size for fragment shuffling (-1 for global)
• nb_epochs: Number of epochs (None for infinite)

fragment_loading · FragmentLoadingConfig · required

Controls how fragments are loaded into memory:

• columns: Schema mapping (use LangASRSchema)
• nb_prefetch: Number of fragments to prefetch
• num_parallel_fragments: Parallel loading threads
• cache: Enable caching of decoded audio

dataset_summary_path · str | None

Path to TSV file with corpus/language hour distribution for weighted sampling.

beta_corpus · float | None

Beta parameter for corpus weighting: weight = (hours/total)^beta.

beta_language · float | None

Beta parameter for language weighting within corpus.

sync_mode · SyncMode · default: "SyncMode.UNTIL_FIRST"

Synchronization mode for distributed training:

• UNTIL_FIRST: Stop when first worker finishes (training)
• UNTIL_LAST: Stop when last worker finishes (validation)

sync_batches · bool · default: "True"

Whether to synchronize batch sizes across workers.

Methods

create_raw_data_pipeline

def create_raw_data_pipeline(
    split: str,
    gangs: Gangs
) -> DataPipelineBuilder

Creates the raw data pipeline for reading parquet files.

split · str · required

Split name (e.g., “train”, “dev”, “test”). Can include corpus filter: “train_librispeech”.

gangs · Gangs · required

Gang configuration for distributed reading.

builder · DataPipelineBuilder

Pipeline builder yielding dictionaries with audio bytes, text, language, and corpus.

Example

from omnilingual_asr.datasets.storage import (
    MixtureParquetStorage,
    MixtureParquetStorageConfig,
    LangASRSchema
)
from fairseq2.data.parquet import (
    FragmentStreamingConfig,
    FragmentLoadingConfig
)
from fairseq2.datasets import SyncMode

# Configure storage
config = MixtureParquetStorageConfig(
    fragment_streaming=FragmentStreamingConfig(
        parquet_path="",
        seed=42,
        fragment_shuffle_window=-1,
        nb_epochs=None
    ),
    fragment_loading=FragmentLoadingConfig(
        columns=LangASRSchema(),
        cache=True
    ),
    dataset_summary_path="stats.tsv",
    beta_corpus=0.5,
    beta_language=0.5,
    sync_mode=SyncMode.UNTIL_FIRST
)

# Create storage
storage = MixtureParquetStorage(
    path=Path("data/asr_parquet"),
    config=config
)

# Create pipeline
builder = storage.create_raw_data_pipeline(
    split="train",
    gangs=gangs
)

AsrTask

ASR preprocessing pipeline including audio filtering, tokenization, and batching.

Constructor

AsrTask(config: AsrTaskConfig)

config · AsrTaskConfig · required

Task configuration for preprocessing pipeline.

Configuration: AsrTaskConfig

Audio Processing

min_audio_len · int · default: "1"

Minimum audio sequence length (in samples). Shorter audio is filtered out.

max_audio_len · int · default: "800000"

Maximum audio sequence length (~50s at 16kHz). Longer audio is filtered out.

normalize_audio · bool · default: "False"

Whether to normalize audio to zero mean and unit variance.

use_fbank · bool · default: "False"

Whether to use filterbank features instead of raw waveforms.

SpecAugment

spec_aug_p · float | None · default: "None"

Probability of applying SpecAugment per sample.

spec_aug_freq_mask_param · int · default: "80"

Maximum frequency mask length for SpecAugment.

spec_aug_time_mask_param · int · default: "80"

Maximum time mask length for SpecAugment.

Text Processing

filter_long_text_threshold · int | None · default: "None"

Maximum text length in tokens. Longer sequences are filtered out.

remove_unknown · bool · default: "False"

Whether to remove unknown tokens from text in-place.

min_samples_per_char · int · default: "160"

Minimum audio samples per character. Samples with faster speech are filtered out.

Batching

batching_strategy · BatchingStrategy · default: "LENGTH"

Batching strategy:

• LENGTH: Dynamic batching by total elements (recommended)
• STATIC: Fixed batch size

batch_size · int · default: "8"

Batch size for STATIC batching strategy.

max_num_elements · int · default: "3200000"

Maximum total elements per batch for LENGTH strategy.

num_seqs_multiple_of · int · default: "8"

Batch size must be multiple of this value (for hardware optimization).

drop_remainder · bool · default: "False"

Whether to drop last incomplete batch.

Pipeline Settings

example_shuffle_window · int · default: "0"

Sliding window size for shuffling examples before batching.

batch_shuffle_window · int · default: "1000"

Sliding window size for shuffling batches.

num_prefetch · int · default: "4"

Number of batches to prefetch in background.

npc · int · default: "10"

Number of parallel calls for data pipeline operations.

Methods

apply_processing_pipeline

def apply_processing_pipeline(
    builder: DataPipelineBuilder,
    gangs: Gangs,
    tokenizer: Tokenizer,
    dtype: torch.dtype
) -> DataPipelineBuilder

Applies the complete ASR preprocessing pipeline.

builder · DataPipelineBuilder · required

Input pipeline builder (typically from storage layer).

gangs · Gangs · required

Gang configuration.

tokenizer · Tokenizer · required

Tokenizer for text encoding.

dtype · torch.dtype · required

Data type for audio tensors.

builder · DataPipelineBuilder

Pipeline builder yielding Seq2SeqBatch objects.

Pipeline Stages

The ASR task pipeline processes data in the following order:

1. Audio Filtering: Filter by length (min_audio_len, max_audio_len)
2. Example Shuffling: Shuffle before batching (example_shuffle_window)
3. Text Tokenization: Encode text with tokenizer
4. Text Filtering: Filter empty text, unknown sequences, long text
5. Batching: Bucket by audio length or static batch size
6. Batch Shuffling: Shuffle batches (batch_shuffle_window)
7. Audio Decoding: Decode audio bytes to waveforms
8. Audio Processing: Normalize, convert to mono, optionally apply SpecAugment
9. Feature Extraction: Extract fbank features (if use_fbank=True)
10. Collation: Collate into padded batches
11. Prefetching: Prefetch batches in background
12. Seq2SeqBatch Conversion: Convert to final batch format

Example

from omnilingual_asr.datasets.tasks import AsrTask, AsrTaskConfig
import torch

# Configure task
config = AsrTaskConfig(
    min_audio_len=8000,
    max_audio_len=800_000,
    normalize_audio=True,
    spec_aug_p=0.5,  # 50% probability of SpecAugment
    spec_aug_freq_mask_param=27,
    spec_aug_time_mask_param=100,
    max_num_elements=3_200_000,
    num_seqs_multiple_of=8,
    example_shuffle_window=1000,
    batch_shuffle_window=100,
    num_prefetch=4
)

# Create task
task = AsrTask(config)

# Apply to storage pipeline
builder = storage.create_raw_data_pipeline("train", gangs)
builder = task.apply_processing_pipeline(
    builder,
    gangs=gangs,
    tokenizer=tokenizer,
    dtype=torch.bfloat16
)

# Finalize pipeline
pipeline = builder.and_return()

Audio Preprocessing

Audio Decoding

Audio bytes are decoded using fairseq2’s AudioDecoder:

from fairseq2.data.audio import AudioDecoder

audio_decoder = AudioDecoder(dtype=torch.float32)
# Decodes to dict: {"waveform": Tensor, "sample_rate": int}

Normalization

def apply_audio_normalization(waveform: Tensor) -> Tensor:
    """Normalize to zero mean and unit variance."""
    return layer_norm(waveform, waveform.shape)

Mono Conversion

def convert_to_mono(waveform: Tensor) -> Tensor:
    """Convert multi-channel to mono by averaging."""
    if waveform.dim() == 2:
        waveform = waveform.mean(dim=1)
    return waveform

SpecAugment

Applied with probability spec_aug_p:

1. Convert waveform to spectrogram
2. Apply frequency masking (random mask of length up to spec_aug_freq_mask_param)
3. Apply time masking (random mask of length up to spec_aug_time_mask_param)
4. Convert back to waveform

Filterbank Features

If use_fbank=True:

from fairseq2.data.audio import WaveformToFbankConverter

fbank_converter = WaveformToFbankConverter(
    num_mel_bins=80,
    waveform_scale=2**15,
    channel_last=True,
    standardize=True,
    dtype=dtype
)

Schema Definitions

LangASRSchema

Column mapping for parquet datasets:

@dataclass
class LangASRSchema(NamedColumns):
    audio: str = "audio_bytes"      # Audio data column
    length: str = "audio_size"       # Audio length column
    text: str = "text"               # Transcription column
    split: str = "split"             # Split column
    lang: str = "language"           # Language column
    corpus: str = "corpus"           # Corpus column

Complete Example: Training Pipeline

from pathlib import Path
import torch
from fairseq2.data.tokenizers import load_tokenizer
from fairseq2.gang import FakeGang
from fairseq2.datasets import SyncMode

from omnilingual_asr.datasets.impl import MixtureParquetAsrDataset
from omnilingual_asr.datasets.storage import (
    MixtureParquetStorageConfig,
    LangASRSchema
)
from omnilingual_asr.datasets.tasks import AsrTaskConfig

# Storage configuration
storage_config = MixtureParquetStorageConfig(
    dataset_summary_path="stats.tsv",
    beta_corpus=0.5,
    beta_language=0.5,
    sync_mode=SyncMode.UNTIL_FIRST,
    sync_batches=True
)

# Task configuration
task_config = AsrTaskConfig(
    normalize_audio=True,
    spec_aug_p=0.5,
    max_num_elements=3_200_000,
    example_shuffle_window=1000,
    batch_shuffle_window=100
)

# Create dataset and reader
dataset = MixtureParquetAsrDataset.from_path(Path("data/asr"))
tokenizer = load_tokenizer("omniASR_tokenizer_v1")

reader = dataset.create_reader(
    split="train",
    tokenizer=tokenizer,
    gangs=FakeGang(device=torch.device("cuda")),
    dtype=torch.bfloat16,
    num_accumulate=1,
    storage_config=storage_config,
    task_config=task_config
)

# Training loop
for batches in reader:
    for batch in batches:
        # batch.source_seqs: [B, T] audio
        # batch.target_seqs: [B, S] text tokens
        loss = model(batch)
        loss.backward()

Source References

• MixtureParquetStorage: src/omnilingual_asr/datasets/storage/mixture_parquet_storage.py:133
• AsrTask: src/omnilingual_asr/datasets/tasks/asr_task.py:140
• Audio utilities: src/omnilingual_asr/datasets/utils/audio.py