"""Configuration and utility classes for classes for Dynamic Chunk Training, as

often used for the training of streaming-capable models in speech recognition.

The definition of Dynamic Chunk Training is based on that of the following

paper, though a lot of the literature refers to the same definition:

https://arxiv.org/abs/2012.05481

Authors

* Sylvain de Langen 2023

"""

from dataclasses import dataclass

from typing import Optional

import torch

import speechbrain as sb

# NOTE: this configuration object is intended to be relatively specific to

# Dynamic Chunk Training; if you want to implement a different similar type of

# chunking different from that, you should consider using a different object.

@dataclass

class DynChunkTrainConfig:

"""Dynamic Chunk Training configuration object for use with transformers,

often in ASR for streaming.

This object may be used both to configure masking at training time and for

run-time configuration of DynChunkTrain-ready models.

"""

chunk_size: int

"""Size in frames of a single chunk, always `>0`.

If chunkwise streaming should be disabled at some point, pass an optional

streaming config parameter."""

left_context_size: Optional[int] = None

"""Number of *chunks* (not frames) visible to the left, always `>=0`.

If zero, then chunks can never attend to any past chunk.

If `None`, the left context is infinite (but use

`.is_infinite_left_context` for such a check)."""

def is_infinite_left_context(self) -> bool:

"""Returns true if the left context is infinite (i.e. any chunk can

attend to any past frame).

"""

return self.left_context_size is None

def left_context_size_frames(self) -> Optional[int]:

"""Returns the number of left context *frames* (not chunks).

If ``None``, the left context is infinite.

See also the ``left_context_size`` field.

"""

if self.left_context_size is None:

return None

return self.chunk_size * self.left_context_size

@dataclass

class DynChunkTrainConfigRandomSampler:

"""Helper class to generate a DynChunkTrainConfig at runtime depending on the current

stage.

Example

-------

>>> from speechbrain.core import Stage

>>> from speechbrain.utils.dynamic_chunk_training import DynChunkTrainConfig

>>> from speechbrain.utils.dynamic_chunk_training import (

... DynChunkTrainConfigRandomSampler,

... )

>>> # for the purpose of this example, we test a scenario with a 100%

>>> # chance of the (24, None) scenario to occur

>>> sampler = DynChunkTrainConfigRandomSampler(

... chunkwise_prob=1.0,

... chunk_size_min=24,

... chunk_size_max=24,

... limited_left_context_prob=0.0,

... left_context_chunks_min=16,

... left_context_chunks_max=16,

... test_config=DynChunkTrainConfig(32, 16),

... valid_config=None,

... )

>>> one_train_config = sampler(Stage.TRAIN)

>>> one_train_config

DynChunkTrainConfig(chunk_size=24, left_context_size=None)

>>> one_train_config.is_infinite_left_context()

True

>>> sampler(Stage.TEST)

DynChunkTrainConfig(chunk_size=32, left_context_size=16)

"""

chunkwise_prob: float

"""When sampling (during `Stage.TRAIN`), the probability that a finite chunk

size will be used.

In the other case, any chunk can attend to the full past and future

context."""

chunk_size_min: int

"""When sampling a random chunk size, the minimum chunk size that can be

picked."""

chunk_size_max: int

"""When sampling a random chunk size, the maximum chunk size that can be

picked."""

limited_left_context_prob: float

"""When sampling a random chunk size, the probability that the left context

will be limited.

In the other case, any chunk can attend to the full past context."""

left_context_chunks_min: int

"""When sampling a random left context size, the minimum number of left

context chunks that can be picked."""

left_context_chunks_max: int

"""When sampling a random left context size, the maximum number of left

context chunks that can be picked."""

test_config: Optional[DynChunkTrainConfig] = None

"""The configuration that should be used for `Stage.TEST`.

When `None`, evaluation is done with full context (i.e. non-streaming)."""

valid_config: Optional[DynChunkTrainConfig] = None

"""The configuration that should be used for `Stage.VALID`.

When `None`, evaluation is done with full context (i.e. non-streaming)."""

def _sample_bool(self, prob):

"""Samples a random boolean with a probability, in a way that depends on

PyTorch's RNG seed.

Arguments

---------

prob : float

Probability (0..1) to return True (False otherwise).

Returns

-------

The sampled boolean

"""

return torch.rand((1,)).item() < prob

def __call__(self, stage):

"""In training stage, samples a random DynChunkTrain configuration.

During validation or testing, returns the relevant configuration.

Arguments

---------

stage : speechbrain.core.Stage

Current stage of training or evaluation.

In training mode, a random DynChunkTrainConfig will be sampled

according to the specified probabilities and ranges.

During evaluation, the relevant DynChunkTrainConfig attribute will

be picked.

Returns

-------

The appropriate configuration

"""

if stage == sb.core.Stage.TRAIN:

# When training for streaming, for each batch, we have a

# `dynamic_chunk_prob` probability of sampling a chunk size

# between `dynamic_chunk_min` and `_max`, otherwise output

# frames can see anywhere in the future.

if self._sample_bool(self.chunkwise_prob):

chunk_size = torch.randint(

self.chunk_size_min,

self.chunk_size_max + 1,

(1,),

).item()

if self._sample_bool(self.limited_left_context_prob):

left_context_chunks = torch.randint(

self.left_context_chunks_min,

self.left_context_chunks_max + 1,

(1,),

).item()

else:

left_context_chunks = None

return DynChunkTrainConfig(chunk_size, left_context_chunks)

return None

elif stage == sb.core.Stage.TEST:

return self.test_config

elif stage == sb.core.Stage.VALID:

return self.valid_config

else:

raise AttributeError(f"Unsupported stage found {stage}")