# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Lhotse CutSet utilities and Parquet manifest support for NeMo.""" import io import logging import random import re import warnings from copy import deepcopy from functools import partial from itertools import repeat from pathlib import Path from typing import KeysView, List, Mapping, Sequence, Tuple, Union import numpy as np import omegaconf import soundfile as sf from lhotse import CutSet, Features, MonoCut, Recording, SupervisionSegment from lhotse.array import Array, TemporalArray from lhotse.cut import Cut, MixedCut, PaddingCut from lhotse.lazy import LazyIteratorChain from lhotse.serialization import load_yaml from lhotse.utils import fastcopy from omegaconf import DictConfig, ListConfig, OmegaConf from nemo.collections.common.data.lhotse.nemo_adapters import ( LazyNeMoIterator, LazyNeMoTarredIterator, LazyParquetIterator, expand_sharded_filepaths, ) from nemo.collections.common.data.lhotse.text_adapters import ( AudioTurn, LhotseTextAdapter, LhotseTextJsonlAdapter, LhotseTextPairAdapter, NeMoMultimodalConversation, NeMoMultimodalConversationJsonlAdapter, NeMoMultimodalConversationShareGPTJsonlAdapter, NeMoMultimodalConversationShareGPTWebdatasetAdapter, NeMoSFTJsonlAdapter, TextTurn, ) from nemo.collections.common.parts.preprocessing.manifest import get_full_path def temperature_reweighting(weights: List[Union[float, int]], temperature: float = 1.0) -> List[float]: """ Apply temperature scaling to dataset weights and normalize. Formula: normalized_weight_i = (w_i ^ temperature) / sum(w_j ^ temperature) Args: weights: List of dataset weights (can be hours, sample counts, or probabilities). Values can be any positive float/int, not limited to [0, 1]. temperature: Scaling factor. - 1.0: preserves original weight ratios - 0.0: equalizes all weights (w^0 = 1) - <1.0: oversamples smaller datasets - >1.0: amplifies weight differences Returns: Normalized weights that sum to 1.0 Example: >>> temperature_reweighting([197, 2159], temperature=1.0) # hours [0.0836, 0.9164] # preserves ratio >>> temperature_reweighting([197, 2159], temperature=0.0) # equalize [0.5, 0.5] """ if len(weights) == 0: return [] weights = np.asarray(weights) if np.any(weights <= 0): raise ValueError(f"All weights must be positive (> 0), got: {weights.tolist()}") weights = weights**temperature return (weights / weights.sum()).tolist() def validate_and_standardize_reweight_temperature(config: Union[DictConfig, dict], propagate_attrs: dict) -> None: """ Validate and standardize reweight_temperature in propagate_attrs. Accepted formats: - Scalar (int/float): broadcast to all nesting levels (warning logged). - List: length must exactly match the input_cfg nesting depth. Raises: ValueError: If list length does not match the nesting depth. """ if propagate_attrs["reweight_temperature"] is None: return expected_length = count_input_cfg_levels(config) reweight_temp = propagate_attrs["reweight_temperature"] if isinstance(reweight_temp, (int, float)): propagate_attrs["reweight_temperature"] = [float(reweight_temp)] * expected_length logging.warning( f"reweight_temperature is a scalar ({reweight_temp}), broadcasting to all {expected_length} levels. " f"Expanded to: {propagate_attrs['reweight_temperature']}" ) else: reweight_temp = list(reweight_temp) if len(reweight_temp) != expected_length: raise ValueError( f"reweight_temperature list length ({len(reweight_temp)}) does not match " f"the input_cfg nesting depth ({expected_length}). " f"Provide exactly {expected_length} values (one per nesting level), " f"or use a scalar to apply the same temperature to all levels." ) propagate_attrs["reweight_temperature"] = reweight_temp def read_cutset_from_config(config: Union[DictConfig, dict]) -> Tuple[CutSet, bool]: """ Reads NeMo configuration and creates a CutSet either from Lhotse or NeMo manifests. Returns a tuple of ``CutSet`` and a boolean indicating whether the data is tarred (True) or not (False). """ # First, check if the dataset is specified in the new configuration format and use it if possible. if not isinstance(config, DictConfig): config = DictConfig(config) if config.get("input_cfg") is not None: cuts, is_tarred = read_dataset_config(config) else: # Now, we'll figure out if we should read Lhotse manifest or NeMo manifest. use_nemo_manifest = all(config.get(opt) is None for opt in ("cuts_path", "shar_path")) if use_nemo_manifest: if config.get("manifest_filepath") is None: raise IncompleteConfigError("You must specify either: manifest_filepath, cuts_path, or shar_path") cuts, is_tarred = read_nemo_manifest(config) else: cuts, is_tarred = read_lhotse_manifest(config) return cuts, is_tarred class IncompleteConfigError(RuntimeError): """Placeholder for an error raised.""" pass KNOWN_DATA_CONFIG_TYPES = {} def get_known_config_data_types() -> KeysView[str]: """ Return the names of all registered data type parsers. Example: >>> get_known_config_data_types() ["nemo", "nemo_tarred", "lhotse", ...] """ return KNOWN_DATA_CONFIG_TYPES.keys() def get_parser_fn(data_type_name: str): """ Return the parsing function for a given data type name. Parsing function reads a dataloading config and returns a tuple of lhotse ``CutSet`` and boolean indicating whether we should use iterable dataset (True) or map dataset (False) mechanism ("is tarred"). """ return KNOWN_DATA_CONFIG_TYPES[data_type_name] def data_type_parser(name: Union[str, list[str]]): """ Decorator used to register data type parser functions. Parsing function reads a dataloading config and returns a tuple of lhotse ``CutSet`` and boolean indicating whether we should use iterable dataset (True) or map dataset (False) mechanism ("is tarred"). Example: >>> @data_type_parser("my_new_format") ... def my_new_format(config): ... return CutSet(read_my_format(**config)), True ... ... fn = get_parser_fn("my_new_format") ... cuts, is_tarred = fn({"my_arg_0": ..., "my_arg_1": ..., ...}) """ def _decorator(fn): global KNOWN_DATA_CONFIG_TYPES if isinstance(name, str): KNOWN_DATA_CONFIG_TYPES[name] = fn else: for n in name: KNOWN_DATA_CONFIG_TYPES[n] = fn return fn return _decorator def read_dataset_config(config) -> tuple[CutSet, bool]: """ Input configuration format examples. Example 1. Combine two datasets with equal weights and attach custom metadata in ``tags`` to each cut:: input_cfg: - type: nemo_tarred manifest_filepath: /path/to/manifest__OP_0..512_CL_.json tarred_audio_filepath: /path/to/tarred_audio/audio__OP_0..512_CL_.tar weight: 0.5 tags: lang: en some_metadata: some_value - type: nemo_tarred manifest_filepath: /path/to/manifest__OP_0..512_CL_.json tarred_audio_filepath: /path/to/tarred_audio/audio__OP_0..512_CL_.tar weight: 0.5 tags: lang: pl some_metadata: some_value Example 2. Combine multiple (4) datasets, with 2 corresponding to different tasks (ASR, AST). There are two levels of weights: per task (outer) and per dataset (inner). The final weight is the product of outer and inner weight:: input_cfg: - type: group weight: 0.7 tags: task: asr input_cfg: - type: nemo_tarred manifest_filepath: /path/to/asr1/manifest__OP_0..512_CL_.json tarred_audio_filepath: /path/to/tarred_audio/asr1/audio__OP_0..512_CL_.tar weight: 0.6 tags: lang: en some_metadata: some_value - type: nemo_tarred manifest_filepath: /path/to/asr2/manifest__OP_0..512_CL_.json tarred_audio_filepath: /path/to/asr2/tarred_audio/audio__OP_0..512_CL_.tar weight: 0.4 tags: lang: pl some_metadata: some_value - type: group weight: 0.3 tags: task: ast input_cfg: - type: nemo_tarred manifest_filepath: /path/to/ast1/manifest__OP_0..512_CL_.json tarred_audio_filepath: /path/to/ast1/tarred_audio/audio__OP_0..512_CL_.tar weight: 0.2 tags: src_lang: en tgt_lang: pl - type: nemo_tarred manifest_filepath: /path/to/ast2/manifest__OP_0..512_CL_.json tarred_audio_filepath: /path/to/ast2/tarred_audio/audio__OP_0..512_CL_.tar weight: 0.8 tags: src_lang: pl tgt_lang: en """ propagate_attrs = { "shuffle": config.get("shuffle", False), "shard_seed": config.get("shard_seed", "trng"), "text_field": config.get("text_field", "text"), "lang_field": config.get("lang_field", "lang"), "metadata_only": config.get("metadata_only", False), "force_finite": config.get("force_finite", False), "max_open_streams": config.get("max_open_streams", None), "audio_locator_tag": config.get("audio_locator_tag", None), "token_equivalent_duration": config.get("token_equivalent_duration", None), "skip_missing_manifest_entries": config.get("skip_missing_manifest_entries", False), "force_map_dataset": config.get("force_map_dataset", False), "force_iterable_dataset": config.get("force_iterable_dataset", False), "slice_length": config.get("slice_length", None), # Temperature for re-weighting datasets. 1 is a neutral value. Lower temperature over-samples smaller datasets, and vice versa. "reweight_temperature": config.get("reweight_temperature", None), } validate_and_standardize_reweight_temperature(config, propagate_attrs) cuts, is_tarred = parse_and_combine_datasets(config.input_cfg, propagate_attrs=propagate_attrs) return cuts, is_tarred def parse_group(grp_cfg: DictConfig, propagate_attrs: dict) -> [CutSet, bool]: """Parse a group configuration, potentially combining multiple datasets.""" assert grp_cfg.type in get_known_config_data_types(), f"Unknown item type in dataset config list: {grp_cfg.type=}" # Note: Text data types will return is_tarred=True. # We choose to treat text as-if it was tarred, which tends to be more # efficient as it moves the text file iteration into dataloading subprocess. if grp_cfg.type != "group": parser_fn = get_parser_fn(grp_cfg.type) cuts, is_tarred = parser_fn(grp_cfg) else: cuts, is_tarred = parse_and_combine_datasets( grp_cfg.input_cfg, propagate_attrs=propagate_attrs, ) # Attach extra tags to every utterance dynamically, if provided. if (extra_tags := grp_cfg.get("tags")) is not None: cuts = cuts.map(partial(attach_tags, tags=extra_tags), apply_fn=None) return cuts, is_tarred @data_type_parser("txt") def read_txt_paths(config: DictConfig) -> tuple[CutSet, bool]: """ Read paths to text files and create a CutSet. """ cuts = CutSet( LhotseTextAdapter( paths=config.paths, language=config.language, shuffle_shards=config.shuffle, shard_seed=config.shard_seed, ) ) if not config.get("force_finite", False): cuts = cuts.repeat(preserve_id=True) return cuts, True @data_type_parser("txt_jsonl") def read_txt_jsonl_paths(config: DictConfig) -> tuple[CutSet, bool]: """Read paths to text files in JSONL format and create a CutSet. This parser reads JSONL files where each line contains a JSON object with text fields. The text_field parameter specifies which field in the JSON object contains the text to be used. """ cuts = CutSet( LhotseTextJsonlAdapter( paths=config.paths, language=config.language, text_field=config.text_field, shuffle_shards=config.shuffle, shard_seed=config.shard_seed, ) ) if not config.get("force_finite", False): cuts = cuts.repeat(preserve_id=True) return cuts, True @data_type_parser("txt_pair") def read_txt_pair_paths(config: DictConfig) -> tuple[CutSet, bool]: """Read paths to source and target text files and create a CutSet.""" cuts = CutSet( LhotseTextPairAdapter( source_paths=config.source_paths, target_paths=config.target_paths, source_language=config.get("source_language"), target_language=config.get("target_language"), questions_path=config.get("questions_path"), questions_language=config.get("questions_language"), shuffle_shards=config.shuffle, shard_seed=config.shard_seed, ) ) if not config.get("force_finite", False): cuts = cuts.repeat(preserve_id=True) return cuts, True @data_type_parser("nemo_sft_jsonl") def read_nemo_sft_jsonl(config: DictConfig) -> tuple[CutSet, bool]: """Read paths to Nemo SFT JSONL files and create a CutSet.""" cuts = CutSet( NeMoSFTJsonlAdapter( paths=config.paths, language=config.get("language"), shuffle_shards=config.shuffle, shard_seed=config.shard_seed, ) ) if not config.get("force_finite", False): cuts = cuts.repeat(preserve_id=True) return cuts, True @data_type_parser("multimodal_conversation") def read_multimodal_conversation_jsonl(config: DictConfig) -> tuple[CutSet, bool]: """Read paths to multimodal conversation JSONL files and create a CutSet.""" cuts = CutSet( NeMoMultimodalConversationJsonlAdapter( manifest_filepath=config.manifest_filepath, tarred_audio_filepaths=config.get("tarred_audio_filepaths"), audio_locator_tag=config.audio_locator_tag, token_equivalent_duration=config.get("token_equivalent_duration"), shuffle_shards=config.shuffle, shard_seed=config.shard_seed, system_prompt=config.get("tags", {}).get("system_prompt"), context=config.get("tags", {}).get("context"), slice_length=config.get("slice_length"), ) ) if not config.get("force_finite", False): cuts = cuts.repeat() return cuts, True @data_type_parser(["share_gpt"]) def read_share_gpt_as_conversation(config) -> tuple[CutSet, bool]: """Read paths to ShareGPT JSONL files and create a CutSet of NeMoMultimodalConversation.""" cuts = CutSet( NeMoMultimodalConversationShareGPTJsonlAdapter( manifest_filepath=config.manifest_filepath, tarred_audio_filepaths=config.get("tarred_audio_filepaths"), audio_locator_tag=config.audio_locator_tag, audio_placeholders=config.audio_placeholders, audio_root=config.get("audio_root"), token_equivalent_duration=config.get("token_equivalent_duration"), shuffle_shards=config.shuffle, shard_seed=config.shard_seed, slice_length=config.get("slice_length"), ) ) if not config.get("force_finite", False): cuts = cuts.repeat(preserve_id=True) return cuts, True @data_type_parser(["share_gpt_webdataset"]) def read_share_gpt_webdataset_as_conversation(config) -> tuple[CutSet, bool]: """Read ShareGPT conversations from WebDataset tar archives.""" cuts = CutSet( NeMoMultimodalConversationShareGPTWebdatasetAdapter( data_dir=config.data_dir, audio_locator_tag=config.audio_locator_tag, audio_placeholders=config.get("audio_placeholders"), token_equivalent_duration=config.get("token_equivalent_duration"), shuffle_shards=config.shuffle, shard_seed=config.shard_seed, ) ) # When force_finite is False (default), repeat the dataset infinitely so that # the dataloader never runs out of data; the trainer controls epoch boundaries. if not config.get("force_finite", False): cuts = cuts.repeat(preserve_id=True) return cuts, True def _resolve_shar_inputs(path: Union[str, Path], only_metadata: bool) -> dict: if only_metadata: return dict(fields={"cuts": sorted(Path(path).glob("cuts.*"))}) else: return dict(in_dir=path) def attach_tags(cut, tags: dict): """Attach extra tags to a cut dynamically.""" for key, val in tags.items(): setattr(cut, key, val) return cut def count_input_cfg_levels(config: Union[DictConfig, dict]) -> int: """ Compute the maximum nesting depth of 'input_cfg' keys in the configuration. Each 'input_cfg' represents one level of nesting that consumes one temperature value from reweight_temperature. Since sibling groups at the same level share the same temperature (due to propagate_attrs.copy()), we count max depth, not total occurrences. String/Path values for ``input_cfg`` are treated as file references (mirroring :func:`parse_and_combine_datasets`) and loaded so that nested ``input_cfg`` keys inside those files are counted. If the file is not found (e.g. the path contains unresolved OmegaConf interpolations such as ``${oc.env:MANIFEST_ROOT}``), it is conservatively counted as one additional level. All other I/O or parsing errors propagate immediately. Args: config: Configuration dictionary that may contain nested 'input_cfg' keys. Returns: Maximum nesting depth of 'input_cfg' keys. Example: >>> config = { ... "input_cfg": [ ... {"type": "group", "input_cfg": [{"type": "nemo"}]}, ... {"type": "group", "input_cfg": [{"type": "nemo"}]}, ... ] ... } >>> count_input_cfg_levels(config) 2 """ _cache: dict[str, object] = {} def _resolve_if_path(val): """If *val* is a string/Path, load the YAML file it points to. Raises on I/O or parse errors except ``FileNotFoundError``, which is expected when the path contains OmegaConf interpolations (e.g. ``${oc.env:MANIFEST_ROOT}/file.yaml``) that raw ``yaml.load`` returns as literal strings. ``parse_and_combine_datasets`` resolves them at runtime via ``OmegaConf.create()``. """ if isinstance(val, (str, Path)): key = str(val) if key not in _cache: try: _cache[key] = load_yaml(key) except FileNotFoundError: logging.debug("count_input_cfg_levels: could not load %r, treating as leaf", key) _cache[key] = val return _cache[key] return val def _max_depth(obj) -> int: if isinstance(obj, (dict, DictConfig)): depths = [] for key, val in obj.items(): if key == "input_cfg": resolved = _resolve_if_path(val) depths.append(1 + _max_depth(resolved)) else: depths.append(_max_depth(val)) return max(depths, default=0) elif isinstance(obj, (list, ListConfig)): # For lists, find the max depth across all items (siblings) return max((_max_depth(item) for item in obj), default=0) return 0 return _max_depth(config) @data_type_parser("group") def parse_and_combine_datasets( config_list: Union[list[DictConfig], ListConfig, str, Path], propagate_attrs: dict ) -> tuple[CutSet, bool]: """Parse a list of dataset configurations, potentially combining multiple datasets.""" cuts = [] weights = [] tarred_status = [] # Extract the temperature for re-weighting datasets. if not propagate_attrs["reweight_temperature"]: temperature = 1.0 next_temperatures = None else: temperature, *next_temperatures = propagate_attrs["reweight_temperature"] propagate_attrs["reweight_temperature"] = next_temperatures if isinstance(config_list, (str, Path)): # Resolve local filepath /path/to/input_cfg.yaml or # remote url s3://bucket/path/to/input_cfg.yaml into config contents if needed. config_list = OmegaConf.create(load_yaml(config_list)) assert len(config_list) > 0, "Empty group in dataset config list." for item in config_list: # Check if we have any attributes that are propagated downwards to each item in the group. # If a key already exists in the item, it takes precedence (we will not overwrite); # otherwise we will assign it. # We also update propagate_atts for the next sub-groups based on what's present in this group next_propagate_attrs = propagate_attrs.copy() for k, v in propagate_attrs.items(): if k not in item: item[k] = v else: next_propagate_attrs[k] = item[k] # Load the item (which may also be another group) as a CutSet. item_cuts, item_is_tarred = parse_group(item, next_propagate_attrs) cuts.append(item_cuts) tarred_status.append(item_is_tarred) if (w := item.get("weight")) is not None: weights.append(w) all_same_tarred_status = all(t == tarred_status[0] for t in tarred_status) if not all_same_tarred_status: if propagate_attrs["force_map_dataset"] or propagate_attrs["force_iterable_dataset"]: logging.warning( f"Not all datasets in the group have the same tarred status, using provided force_map_dataset " f"({propagate_attrs['force_map_dataset']}) and force_iterable_dataset " f"({propagate_attrs['force_iterable_dataset']}) to determine the final tarred status." ) else: raise ValueError( "Mixing tarred and non-tarred datasets is not supported when neither force_map_dataset " "nor force_iterable_dataset is True." ) assert len(weights) == 0 or len(cuts) == len( weights ), "Missing dataset weight. When weighting datasets, every dataset must have a specified weight." if len(cuts) > 1: if not weights: reweights = None else: reweights = temperature_reweighting(weights, temperature=temperature) cuts = mux( *cuts, weights=reweights, max_open_streams=propagate_attrs["max_open_streams"], seed=propagate_attrs["shard_seed"], force_finite=propagate_attrs["force_finite"] or propagate_attrs["metadata_only"], ) else: (cuts,) = cuts return cuts, tarred_status[0] @data_type_parser(["lhotse", "lhotse_shar"]) def read_lhotse_manifest(config) -> tuple[CutSet, bool]: """Read paths to Lhotse manifest files and create a CutSet.""" is_tarred = config.get("shar_path") is not None if is_tarred: # Lhotse Shar is the equivalent of NeMo's native "tarred" dataset. # The combination of shuffle_shards, and repeat causes this to # be an infinite manifest that is internally reshuffled on each epoch. # The parameter ``config.shard_seed`` is used to determine shard shuffling order. Options: # - "trng" means we'll defer setting the seed until the iteration # is triggered, and we'll use system TRNG to get a completely random seed for each worker. # This results in every dataloading worker using full data but in a completely different order. # - "randomized" means we'll defer setting the seed until the iteration # is triggered, and we'll use config.seed to get a pseudo-random seed for each worker. # This results in every dataloading worker using full data but in a completely different order. # Unlike "trng", this is deterministic, and if you resume training, you should change the seed # to observe different data examples than in the previous run. # - integer means we'll set a specific seed in every worker, and data would be duplicated across them. # This is mostly useful for unit testing or debugging. shard_seed = config.get("shard_seed", "trng") metadata_only = config.get("metadata_only", False) force_finite = config.get("force_finite", False) if config.get("cuts_path") is not None: warnings.warn("Note: lhotse.cuts_path will be ignored because lhotse.shar_path was provided.") if isinstance(config.shar_path, (str, Path)): cuts = CutSet.from_shar( **_resolve_shar_inputs(config.shar_path, metadata_only), shuffle_shards=True, seed=shard_seed, slice_length=config.get("slice_length", None), ) if not metadata_only and not force_finite: cuts = cuts.repeat(preserve_id=True) elif isinstance(config.shar_path, Sequence): # Multiple datasets in Lhotse Shar format: we will dynamically multiplex them # with probability approximately proportional to their size cutsets = [] weights = [] for item in config.shar_path: if isinstance(item, (str, Path)): path = item cs = CutSet.from_shar( **_resolve_shar_inputs(path, metadata_only), shuffle_shards=True, seed=shard_seed, slice_length=config.get("slice_length", None), ) weight = len(cs) else: assert isinstance(item, Sequence) and len(item) == 2 and isinstance(item[1], (int, float)), ( "Supported inputs types for config.shar_path are: " "str | list[str] | list[tuple[str, number]] " "where str is a path and number is a mixing weight (it may exceed 1.0). " f"We got: '{item}'" ) path, weight = item cs = CutSet.from_shar( **_resolve_shar_inputs(path, metadata_only), shuffle_shards=True, seed=shard_seed, slice_length=config.get("slice_length", None), ) cutsets.append(cs) weights.append(weight) cuts = mux( *cutsets, weights=weights, max_open_streams=config.get("max_open_streams", None), seed=shard_seed, force_finite=force_finite, ) elif isinstance(config.shar_path, Mapping): fields = {k: expand_sharded_filepaths(v) for k, v in config.shar_path.items()} assert "cuts" in config.shar_path.keys(), ( f"Invalid value for key 'shar_path': a dict was provided, but didn't specify key 'cuts' pointing " f"to the manifests. We got the following: {config.shar_path=}" ) if metadata_only: fields = {"cuts": fields["cuts"]} cuts = CutSet.from_shar( fields=fields, shuffle_shards=True, seed=shard_seed, slice_length=config.get("slice_length", None), ) if not metadata_only and not force_finite: cuts = cuts.repeat(preserve_id=True) else: raise RuntimeError( f"Unexpected value for key 'shar_path'. We support string, list of strings, " f"list of tuples[string,float], and dict[string,list[string]], " f"but got: {type(config.shar_path)=} {config.shar_path=}" ) else: # Regular Lhotse manifest points to individual audio files (like native NeMo manifest). path = config.cuts_path cuts = CutSet.from_file(path).map(partial(resolve_relative_paths, manifest_path=path)) return cuts, is_tarred @data_type_parser(["parquet"]) def read_parquet_manifest(config: DictConfig) -> tuple[CutSet, bool]: """ Read paths to Parquet files and create a CutSet. Config options: - manifest_filepath: path to .parquet file (or list of paths) - audio_field: column name for audio (default: "audio") - text_field: column name for text (default: "text") - duration_field: column name for duration (default: "duration") - lang_field: column name for language (default: "lang") - sampling_rate: default sampling rate if not present (default: 16000) - shuffle: whether to shuffle shards (default: False) - shard_seed: seed for shuffling (default: "trng") """ # 1. Get the path(s) paths = config.manifest_filepath if isinstance(paths, str): paths = [paths] # 2. Extract config options with defaults audio_field = config.get("audio_field", "audio") text_field = config.get("text_field", "text") duration_field = config.get("duration_field", "duration") lang_field = config.get("lang_field", "lang") sampling_rate = config.get("sampling_rate", 16000) # Extract shuffling options (CRITICAL for distributed training) shuffle_shards = config.get("shuffle", False) shard_seed = config.get("shard_seed", "trng") # 3. Create Iterators for each file iterators = [] for path in paths: logging.info(f"Initializing Lhotse Parquet Iterator: '{path}'") adapter = LazyParquetIterator( path=path, audio_field=audio_field, text_field=text_field, duration_field=duration_field, lang_field=lang_field, sampling_rate=sampling_rate, ) iterators.append(adapter) # 4. Chain them together using Lhotse's lazy chaining if len(iterators) == 1: source = iterators[0] else: # This handles shuffling the order of .parquet files for multi-GPU training # as requested by pzelasko source = LazyIteratorChain(*iterators, shuffle_iters=shuffle_shards, seed=shard_seed) # 5. Create the final CutSet cuts = CutSet(source) # 6. Handle infinite looping if requested if not config.get("force_finite", False): cuts = cuts.repeat(preserve_id=True) # Return cuts + is_tarred=True (since it's a stream, we treat it like tarred) return cuts, True def cut_to_conversation( cut: Cut, audio_locator_tag: str, token_equivalent_duration: float ) -> NeMoMultimodalConversation: """ Converts a lhotse Cut into a two-turn NeMoMultimodalConversation, where the user turn contains cut's audio, and assistant turn contains text response from ``cut.supervisions[0].text``. If ``cut`` has a custom field ``context``, it's pre-pended as an extra user text turn before the user's audio turn. """ if isinstance(cut, NeMoMultimodalConversation): return cut turns = [ AudioTurn(cut=cut, role="user", audio_locator_tag=audio_locator_tag, text=cut.supervisions[0].text), TextTurn(value=cut.supervisions[0].text, role="assistant"), ] if hasattr(cut, "context"): turns = [TextTurn(value=cut.context, role="user")] + turns if hasattr(cut, "system_prompt"): turns = [TextTurn(value=cut.system_prompt, role="system")] + turns return NeMoMultimodalConversation( id=cut.id, turns=turns, token_equivalent_duration=token_equivalent_duration, custom=cut.custom, ) @data_type_parser(["s2s_duplex_overlap_as_s2s_duplex"]) def read_s2s_duplex_overlap_as_s2s_duplex(config) -> Tuple[CutSet, bool]: """ Convert a CutSet with overlapping agent/user segments into a standard S2S duplex format. Use Case: This parser is designed for conversational data where agent and user speech can overlap in time (e.g., natural turn-taking with interruptions or backchanneling). The input format stores agent and user segments separately as `agent_segments` and `user_segments` attributes on each cut. This function converts them into a unified timeline of sequential SupervisionSegments, which is the standard format expected by DuplexS2S models. Expected Input Data Format: Each cut should have: - cut.agent_segments: List[Dict] with keys: - "start" (float): Start time in seconds - "end" (float): End time in seconds - "text" (str): Agent's transcription - cut.user_segments: List[Dict] with keys: - "start" (float): Start time in seconds - "end" (float): End time in seconds - "text" (str): User's transcription Example: Input cut with overlapping segments: cut.agent_segments = [ {"start": 0.5, "end": 2.0, "text": "Hello, how can I help?"}, {"start": 3.0, "end": 4.5, "text": "Sure, I can do that."} ] cut.user_segments = [ {"start": 1.8, "end": 3.2, "text": "I need assistance"}, {"start": 4.0, "end": 5.5, "text": "Thank you"} ] Output cut.supervisions (sorted by start time): [ SupervisionSegment(start=0.5, duration=1.5, text="Hello, how can I help?", speaker="agent"), SupervisionSegment(start=1.8, duration=1.4, text="I need assistance", speaker="user"), SupervisionSegment(start=3.0, duration=1.5, text="Sure, I can do that.", speaker="agent"), SupervisionSegment(start=4.0, duration=1.5, "Thank you", speaker="user") ] Args: config: Dictionary containing parser options: - move_agent_text_back_by (float): Time offset to shift agent text back (default: 0). Useful for aligning agent text with earlier audio timing. - filter_samples_starting_with_agent (bool): Whether to remove samples starting with agent (default: False). When True, only keeps samples where the first speaker is a user. - agent_roles (List[str]): Roles considered as agent (default: ["agent", "Assistant", "assistant"]). Returns: Tuple[CutSet, bool]: Converted cuts with unified supervisions, and a flag indicating if the data was tarred. """ move_agent_text_back_by = config.get("move_agent_text_back_by", 0) filter_samples_starting_with_agent = config.get("filter_samples_starting_with_agent", False) agent_roles = config.get("agent_roles", ["agent", "Assistant", "assistant"]) cuts, is_tarred = read_cutset_from_config(config) def filter_cuts_starting_with_agent_fn(cuts: CutSet, agent_roles: Tuple[str, ...]) -> CutSet: """Remove cuts where the first supervision belongs to an agent role.""" def _filter_fn(cut: Cut) -> bool: if not cut.supervisions: return False cut.supervisions = sorted(cut.supervisions, key=lambda s: s.start) return cut.supervisions[0].speaker not in agent_roles return cuts.filter(_filter_fn) def convert_overlap_cut_fn(cut: Cut) -> Cut: """Convert agent/user overlapping segments into sequential SupervisionSegments.""" agent_segments = [ SupervisionSegment( id=cut.id, recording_id=cut.id, start=seg["start"] - move_agent_text_back_by, duration=seg["end"] - seg["start"] + move_agent_text_back_by, text=seg["text"], speaker="agent", ) for seg in cut.agent_segments ] user_segments = [ SupervisionSegment( id=cut.id, recording_id=cut.id, start=seg["start"], duration=seg["end"] - seg["start"], text=seg["text"], speaker="user", ) for seg in cut.user_segments ] cut.supervisions = sorted(agent_segments + user_segments, key=lambda s: s.start) cut.task = "s2s_duplex_overlap_as_s2s_duplex" return cut cuts = cuts.map(convert_overlap_cut_fn) if filter_samples_starting_with_agent: cuts = filter_cuts_starting_with_agent_fn(cuts, tuple(agent_roles)) return cuts, is_tarred @data_type_parser(["lhotse_magpietts_data_as_continuation"]) def read_lhotse_magpietts_data_as_s2s_duplex(config) -> Tuple[CutSet, bool]: """ Convert MagpieTTS dataset cuts into the Duplex S2S format, with optional `context_audio` that can be used as a speaker reference. Args: config: Dictionary containing parser options: - add_extra_end_silence (bool): Whether to add extra silence at the end. - extra_end_silence_range (List[float]): Range of extra silence duration. - max_cer (float): Maximum allowed character error rate. - min_context_speaker_similarity (float): Minimum similarity score. - target_speaker (str, optional): Target speaker filter. - sample_rate (int): Audio sample rate for resampling. Returns: Tuple[CutSet, bool]: Converted cuts and a flag indicating if data was tarred. """ cuts, is_tarred = read_cutset_from_config(config) add_extra_end_sil = config.get("add_extra_end_silence", False) extra_end_silence_range = config.get("extra_end_silence_range", [0.5, 6.0]) sample_rate = config.get("sample_rate", 22050) max_cer = config.get("max_cer", 0.03) min_context_speaker_similarity = config.get("min_context_speaker_similarity", 0.6) target_speaker = config.get("target_speaker", None) keep_flag = "pass" def create_recording_from_array(samples: np.ndarray, sampling_rate: int, recording_id: str) -> Recording: """Convert a numpy array into a Lhotse Recording object.""" with io.BytesIO() as buffer: sf.write(buffer, samples.T, samplerate=sampling_rate, format='WAV') buffer.seek(0) return Recording.from_bytes(buffer.read(), recording_id=recording_id) def convert_cut_fn(cut: Cut) -> Cut: """Convert a single cut into the continuation format.""" orig_agent_sup = deepcopy(cut.supervisions[0]) target_audio_orig_dur = cut.target_audio.duration # Resample audios cut.target_audio = cut.target_audio.resample(sample_rate) cut.context_audio = cut.context_audio.resample(sample_rate) total_duration = cut.target_audio.duration # Prepare MonoCuts cut_target = MonoCut( id=f"{cut.id}_target", start=0.0, duration=total_duration, channel=0, recording=cut.target_audio, supervisions=[], ) zero_audio = np.zeros((1, int(total_duration * sample_rate)), dtype=np.float32) source_recording = create_recording_from_array(zero_audio, sample_rate, recording_id=f"{cut.id}_source") cut_source = MonoCut( id=f"{cut.id}_source", start=0.0, duration=total_duration, channel=0, recording=source_recording, supervisions=[], ) # Save to memory cut_source = cut_source.move_to_memory(audio_format='wav') cut_target = cut_target.move_to_memory(audio_format='wav') # Create user and agent supervisions user_sup = fastcopy(orig_agent_sup, start=0.0, duration=0.08, speaker="user", text="dummy text") agent_sup = fastcopy(orig_agent_sup, start=0.0, duration=target_audio_orig_dur - 0.08, speaker="agent") # Optionally add extra silence if add_extra_end_sil: sil_duration = random.uniform(*extra_end_silence_range) cut_target = cut_target.pad(duration=total_duration + sil_duration, direction="right") cut_source = cut_source.pad(duration=total_duration + sil_duration, direction="right") cut_source = cut_source.to_mono().move_to_memory(audio_format='wav') cut_target = cut_target.to_mono().move_to_memory(audio_format='wav') agent_sup.duration += sil_duration + 1.0 user_sup.duration += sil_duration # Assemble final cut cut_source.supervisions = [user_sup, agent_sup] cut_source.target_audio = cut_target.recording cut_source.duration = cut_target.duration cut_source.context_audio = cut.context_audio cut_source.task = "lhotse_magpietts_data_as_continuation" return cut_source # Filters def filter_cer_fn(cut: Cut) -> bool: return ( len(cut.supervisions) == 0 or not cut.supervisions[0].has_custom("cer") or cut.supervisions[0].cer <= max_cer ) def filter_val_flag_fn(cut: Cut) -> bool: return not cut.has_custom("validation_status") or cut.validation_status == keep_flag def filter_secs_fn(cut: Cut) -> bool: return ( len(cut.supervisions) == 0 or not cut.supervisions[0].has_custom("context_speaker_similarity") or cut.supervisions[0].context_speaker_similarity >= min_context_speaker_similarity ) def filter_target_speaker_fn(cut: Cut) -> bool: return len(cut.supervisions) == 0 or target_speaker is None or target_speaker in cut.supervisions[0].speaker # Apply filters cuts = ( cuts.filter(filter_cer_fn).filter(filter_val_flag_fn).filter(filter_secs_fn).filter(filter_target_speaker_fn) ) # Convert cuts cuts = cuts.map(convert_cut_fn) return cuts, is_tarred @data_type_parser(["s2s_duplex_reverse_role"]) def read_s2s_duplex_reverse_role(config) -> Tuple[CutSet, bool]: """ Reverse the speaker roles and swap the source/target audio streams in a Duplex S2S CutSet. This parser takes an existing conversational dataset and inverts the perspective by swapping the "user" and "agent" supervision labels. It also swaps the primary `recording` (usually source audio) with the `target_audio` to fully simulate the conversation from the opposite participant's point of view. Args: config: Dictionary containing parser options: - agent_roles (List[str], optional): List of role strings to be identified as the agent. Defaults to ["agent", "Agent", "Assistant", "assistant"]. - user_roles (List[str], optional): List of role strings to be identified as the user. Defaults to ["user", "User"]. - target_agent_name (str, optional): The canonical name to assign to former user roles. Defaults to "agent". - target_user_name (str, optional): The canonical name to assign to former agent roles. Defaults to "user". Returns: Tuple[CutSet, bool]: Converted cuts with swapped roles and audio streams, along with a flag indicating if the data was tarred. """ cuts, is_tarred = read_cutset_from_config(config) # Roles coming from config agent_roles = config.get("agent_roles", ["agent", "Agent", "Assistant", "assistant"]) user_roles = config.get("user_roles", ["user", "User"]) # Normalize for robust matching agent_roles_set = {r.lower() for r in agent_roles} user_roles_set = {r.lower() for r in user_roles} # Canonical names you want after swapping target_agent_name = config.get("target_agent_name", "agent") target_user_name = config.get("target_user_name", "user") def swap_speaker(role: str) -> str: """Swap a given role based on the configured user/agent sets.""" if role is None: return role role_l = role.lower() # user -> agent if role_l in user_roles_set: return target_agent_name # agent -> user if role_l in agent_roles_set: return target_user_name # untouched roles (e.g., narrator, system, etc.) return role def convert_cut_fn(cut: Cut) -> Cut: """Convert a single cut by swapping supervisions and audio streams.""" new_cut = deepcopy(cut) # swap supervisions if getattr(new_cut, "supervisions", None): new_sups = [] for s in new_cut.supervisions: s2 = deepcopy(s) s2.speaker = swap_speaker(getattr(s2, "speaker", None)) new_sups.append(s2) new_cut.supervisions = new_sups # swap audio streams old_recording = new_cut.recording old_target_audio = new_cut.target_audio old_rec_id = old_recording.id old_tar_id = old_target_audio.id new_cut.recording = old_target_audio new_cut.target_audio = old_recording # keep duration consistent if hasattr(new_cut, "duration"): new_cut.duration = new_cut.recording.duration # Debug assertions assert new_cut.target_audio.id == old_rec_id, f"{new_cut.id}: recording swap failed" assert new_cut.recording.id == old_tar_id, f"{new_cut.id}: target_audio swap failed" # Optional stronger assertions (object identity) assert new_cut.recording is old_target_audio, f"{new_cut.id}: recording object not swapped" assert new_cut.target_audio is old_recording, f"{new_cut.id}: target_audio object not swapped" new_cut.task = "s2s_duplex_reverse_role" return new_cut cuts = cuts.map(convert_cut_fn) return cuts, is_tarred @data_type_parser(["lhotse_as_conversation"]) def read_lhotse_as_conversation(config) -> tuple[CutSet, bool]: """ Conversion parser for regular ASR data. Intended for converting ASR data in NeMo or Lhotse manifests from Lhotse Cut into NeMoMultimodalConversation. The examples for multimodal LLM are constructed to present an optional "context" field and the acoustic representation as input, and predict "text". Example of original data JSON in NeMo manifest format:: { "audio_filepath": "/path/to/audio.wav", "duration": 3.48, "text": "Of exclusive national competence, let me assure you". "lang": "en", "context": "Transcribe the following:", } Same example in Lhotse manifest format:: { "id": "some-id-0", "recording": { "id": "some-id-0", "sources": [ "type": "path", "source": "/path/to/audio.wav", "channel_ids": [0], ], "sampling_rate": 16000, "duration": 3.48, "num_samples": 55680, }, "supervisions": [ "id": "some-id-0", "start": 0.0, "duration": 3.48, "channel": 0, "text": "Of exclusive national competence, let me assure you". "language": "en", ], "start": 0.0, "duration": 3.48, "channel": 0, "custom": { "context": "Transcribe the following:", } } """ cuts, is_tarred = read_cutset_from_config(config) # Attach extra tags to every utterance dynamically, if provided. # We need to attach them before cuts are converted to conversations. if (extra_tags := config.get("tags")) is not None: cuts = cuts.map(partial(attach_tags, tags=extra_tags), apply_fn=None) cuts = cuts.map( partial( cut_to_conversation, audio_locator_tag=config.audio_locator_tag, token_equivalent_duration=config.token_equivalent_duration, ) ) return cuts, is_tarred def sqa_cut_to_conversation( cut: Cut, audio_locator_tag: str, token_equivalent_duration: float ) -> NeMoMultimodalConversation: """ Converts a lhotse Cut representing a SQA example into a NeMoMultimodalConversation. The ``cut`` is expected to have ``question`` and ``answer`` fields. """ if isinstance(cut, NeMoMultimodalConversation): return cut turns = [ TextTurn(value=cut.question, role="user"), AudioTurn(cut=cut, role="user", audio_locator_tag=audio_locator_tag, text=getattr(cut, "text", "")), TextTurn(value=cut.answer, role="assistant"), ] if hasattr(cut, "context"): turns = [TextTurn(value=cut.context, role="user")] + turns if hasattr(cut, "system_prompt"): turns = [TextTurn(value=cut.system_prompt, role="system")] + turns return NeMoMultimodalConversation( id=cut.id, turns=turns, token_equivalent_duration=token_equivalent_duration, custom=cut.custom, ) @data_type_parser(["sqa_as_conversation"]) def read_sqa_as_conversation(config) -> tuple[CutSet, bool]: """ Conversion parser for old spoken-question-answering data saved in NeMo format. Intended for converting SQA data in NeMo or Lhotse manifests from Lhotse Cut into NeMoMultimodalConversation. The examples for multimodal LLM are constructed to present "question" + acoustic representation as input, and predict "answer". Example of original data JSON in NeMo manifest format:: { "audio_filepath": "/path/to/audio.wav", "duration": 3.48, "text": "Of exclusive national competence, let me assure you". "lang": "en", "question": "What is the nature of the competence described in the context?", "answer": "The context mentions \"exclusive national competence\", indicating that the competence in question is solely within the authority of a nation.", } Same example in Lhotse manifest format:: { "id": "some-id-0", "recording": { "id": "some-id-0", "sources": [ "type": "path", "source": "/path/to/audio.wav", "channel_ids": [0], ], "sampling_rate": 16000, "duration": 3.48, "num_samples": 55680, }, "supervisions": [ "id": "some-id-0", "start": 0.0, "duration": 3.48, "channel": 0, "text": "Of exclusive national competence, let me assure you". "language": "en", ], "start": 0.0, "duration": 3.48, "channel": 0, "custom": { "question": "What is the nature of the competence described in the context?", "answer": "The context mentions \"exclusive national competence\", indicating that the competence in question is solely within the authority of a nation.", } } """ cuts, is_tarred = read_cutset_from_config(config) # Attach extra tags to every utterance dynamically, if provided. # We need to attach them before cuts are converted to conversations. if (extra_tags := config.get("tags")) is not None: cuts = cuts.map(partial(attach_tags, tags=extra_tags), apply_fn=None) cuts = cuts.map( partial( sqa_cut_to_conversation, audio_locator_tag=config.audio_locator_tag, token_equivalent_duration=config.token_equivalent_duration, ) ) return cuts, is_tarred def _strip_timestamps( text: str, _TIMESTAMP_PATTERN=re.compile(r"<\|\d+\|>"), _SPACE_PATTERN=re.compile(r"\s+") ) -> str: """ Strips timestamp tokens from text, e.g. turns: '<|0|> Hey <|3|> <|3|> how <|5|> <|7|> are <|8|> <|8|> <|10|> you? <|12|>' into: 'Hey how are you?' """ # Regexp pattern args are cached compiled patterns (micro-optimization). text = _TIMESTAMP_PATTERN.sub("", text) # strip timestamp tokens if present return _SPACE_PATTERN.sub(" ", text).strip() # strip multi-whitespaces class FailedConversion: pass def s2s_cut_to_conversation( cut: Cut, audio_locator_tag: str, token_equivalent_duration: float, input_roles: Sequence[str] = ("user", "User"), output_roles: Sequence[str] = ("assistant", "Assistant", "agent", "Agent"), strip_timestamp_tokens: bool = True, ) -> NeMoMultimodalConversation: """ Converts a lhotse Cut representing multi-turn speech-to-speech conversation (with multiple supervision segments) into a multi-turn NeMoMultimodalConversation, where the user has AudioTurns and assistant responds in TextTurns. Args: cut: lhotse Cut to convert. audio_locator_tag: special token indicating audio will be inserted in this location in the token sequence. token_equivalent_duration: how much speech duration is counted as one token. input_roles: when supervision.speaker is set to one of these values, we consider it user's turn. output_roles: when supervision.speaker is set to one of these values, we consider it assistant's turn. strip_timestamp_tokens: strips tokens like <|0|>, <|1|>, etc indicating timestamps from the text. """ turn_cuts = cut.trim_to_supervisions(keep_overlapping=False) turns = [] idx = 0 for per_turn_cut in turn_cuts: assert ( len(per_turn_cut.supervisions) >= 1 ), f"Expected at least one supervision per turn, got none in cut {cut.id}" # If len(per_turn_cut.supervisions) > 1, only the first turn is considered for cut creation. # We assume that len(per_turn_cut.supervisions) >= 1 happens because one of the turns # is completely contained within another turn. turn_speaker = per_turn_cut.supervisions[0].speaker turn_text = per_turn_cut.supervisions[0].text if strip_timestamp_tokens: turn_text = _strip_timestamps(turn_text) if len(per_turn_cut.supervisions) > 1: assert per_turn_cut.supervisions[1].text == turn_cuts[idx - 1].supervisions[0].text if turn_speaker in input_roles: turns.append(AudioTurn(cut=per_turn_cut, role="user", audio_locator_tag=audio_locator_tag, text=turn_text)) elif turn_speaker in output_roles: turns.append(TextTurn(value=turn_text, role="assistant")) else: logging.warning(f"Speaker '{turn_speaker}' not found in user or agent roles for cut {cut.id}") return FailedConversion() idx += 1 if hasattr(cut, "context"): turns = [TextTurn(value=cut.context, role="user")] + turns if hasattr(cut, "system_prompt"): turns = [TextTurn(value=cut.system_prompt, role="system")] + turns return NeMoMultimodalConversation( id=cut.id, turns=turns, token_equivalent_duration=token_equivalent_duration, custom=cut.custom, ) @data_type_parser(["s2s_as_conversation"]) def read_s2s_as_conversation(config) -> tuple[CutSet, bool]: cuts, is_tarred = read_cutset_from_config(config) # Attach extra tags to every utterance dynamically, if provided. # We need to attach them before cuts are converted to conversations. if (extra_tags := config.get("tags")) is not None: cuts = cuts.map(partial(attach_tags, tags=extra_tags), apply_fn=None) cuts = cuts.map( partial( s2s_cut_to_conversation, audio_locator_tag=config.audio_locator_tag, token_equivalent_duration=config.token_equivalent_duration, input_roles=config.get("input_roles", ["user", "User"]), output_roles=config.get("output_roles", ["assistant", "Assistant", "agent", "Agent"]), strip_timestamp_tokens=config.get("strip_timestamp_tokens", True), ) ).filter(lambda ex: not isinstance(ex, FailedConversion)) return cuts, is_tarred def create_recording_from_array(samples: np.ndarray, sampling_rate: int, recording_id: str) -> Recording: with io.BytesIO() as buffer: sf.write(buffer, samples.T, samplerate=sampling_rate, format='WAV') buffer.seek(0) return Recording.from_bytes(buffer.read(), recording_id=recording_id) def resolve_relative_paths(cut: Cut, manifest_path: str) -> Cut: """Resolve relative paths in a Cut object to their full paths.""" if isinstance(cut, PaddingCut): return cut if isinstance(cut, MixedCut): for track in cut.tracks: track.cut = resolve_relative_paths(track.cut, manifest_path) return cut def resolve_recording(value): for audio_source in value.sources: if audio_source.type == "file": audio_source.source = get_full_path(audio_source.source, manifest_file=manifest_path) def resolve_array(value): if isinstance(value, TemporalArray): value.array = resolve_array(value.array) else: if value.storage_type in ("numpy_files", "lilcom_files"): abspath = Path( get_full_path(str(Path(value.storage_path) / value.storage_key), manifest_file=manifest_path) ) value.storage_path = str(abspath.parent) value.storage_key = str(abspath.name) elif value.storage_type in ( "kaldiio", "chunked_lilcom_hdf5", "lilcom_chunky", "lilcom_hdf5", "numpy_hdf5", ): value.storage_path = get_full_path(value.storage_path, manifest_file=manifest_path) # ignore others i.e. url, in-memory data, etc. if cut.has_recording: resolve_recording(cut.recording) if cut.has_features: resolve_array(cut.features) if cut.custom is not None: for key, value in cut.custom.items(): if isinstance(value, Recording): resolve_recording(value) elif isinstance(value, (Array, TemporalArray, Features)): resolve_array(value) return cut @data_type_parser(["nemo", "nemo_tarred"]) def read_nemo_manifest(config) -> tuple[CutSet, bool]: """Read NeMo manifest and return a Lhotse CutSet.""" common_kwargs = {} for key in ("text_field", "lang_field", "shuffle", "shard_seed", "extra_fields"): if key in config: if key == "shuffle": common_kwargs["shuffle_shards"] = config[key] else: common_kwargs[key] = config[key] # The option below is to allow a special case of NeMo manifest iteration as Lhotse CutSet # without performing any I/O. NeMo manifests typically don't have sampling_rate information required by Lhotse, # so lhotse has to look up the headers of audio files to fill it on-the-fly. # (this only has an impact on non-tarred data; tarred data is read into memory anyway). # This is useful for utility scripts that iterate metadata and estimate optimal batching settings # and other data statistics. metadata_only = config.get("metadata_only", False) force_finite = config.get("force_finite", False) notar_kwargs = {"metadata_only": metadata_only} is_tarred = config.get("tarred_audio_filepaths") is not None if isinstance(config.manifest_filepath, (str, Path)): if is_tarred and not metadata_only: cuts = CutSet( LazyNeMoTarredIterator( config.manifest_filepath, tar_paths=config.tarred_audio_filepaths, skip_missing_manifest_entries=config.get("skip_missing_manifest_entries", False), slice_length=config.get("slice_length", None), **common_kwargs, ) ) if not force_finite: cuts = cuts.repeat(preserve_id=True) else: cuts = CutSet(LazyNeMoIterator(config.manifest_filepath, **notar_kwargs, **common_kwargs)) else: # Format option 1: # Assume it's [[path1], [path2], ...] (same for tarred_audio_filepaths). # This is the format for multiple NeMo buckets. # Note: we set "weights" here to be proportional to the number of utterances in each data source. # this ensures that we distribute the data from each source uniformly throughout each epoch. # Setting equal weights would exhaust the shorter data sources closer the towards the beginning # of an epoch (or over-sample it in the case of infinite CutSet iteration with .repeat()). # Format option 2: # Assume it's [[path1, weight1], [path2, weight2], ...] (while tarred_audio_filepaths remain unchanged). # Note: this option allows to manually set the weights for multiple datasets. # Format option 3: # i.e., NeMo concatenated dataset # Assume it's [path1, path2, ...] (while tarred_audio_filepaths in the same format). cutsets = [] weights = [] tar_paths = config.tarred_audio_filepaths if is_tarred else repeat((None,)) # Create a stream for each dataset. for manifest_info, tar_path in zip(config.manifest_filepath, tar_paths): if is_tarred and isinstance(tar_path, (list, tuple, ListConfig)): # if it's in option 1 or 2 (tar_path,) = tar_path manifest_path = manifest_info[0] else: # if it's in option 3 manifest_path = manifest_info # First, convert manifest_path[+tar_path] to an iterator. if is_tarred and not metadata_only: nemo_iter = LazyNeMoTarredIterator( manifest_path=manifest_path, tar_paths=tar_path, skip_missing_manifest_entries=config.get("skip_missing_manifest_entries", False), slice_length=config.get("slice_length", None), **common_kwargs, ) else: nemo_iter = LazyNeMoIterator(manifest_path, **notar_kwargs, **common_kwargs) # Then, determine the weight or use one provided if isinstance(manifest_info, str) or len(manifest_info) == 1: weight = len(nemo_iter) else: assert ( isinstance(manifest_info, Sequence) and len(manifest_info) == 2 and isinstance(manifest_info[1], (int, float)) ), ( "Supported inputs types for config.manifest_filepath are: " "str | list[list[str]] | list[tuple[str, number]] " "where str is a path and number is a mixing weight (it may exceed 1.0). " f"We got: '{manifest_info}'" ) weight = manifest_info[1] # [optional] When we have a limit on the number of open streams, # split the manifest to individual shards if applicable. # This helps the multiplexing achieve closer data distribution # to the one desired in spite of the limit. if config.get("max_open_streams") is not None: for subiter in nemo_iter.to_shards(): cutsets.append(CutSet(subiter)) weights.append(weight) else: cutsets.append(CutSet(nemo_iter)) weights.append(weight) cuts = mux( *cutsets, weights=weights, max_open_streams=config.get("max_open_streams"), seed=config.get("shard_seed", "trng"), force_finite=force_finite or metadata_only, ) return cuts, is_tarred @data_type_parser("multi_speaker_simulator") def read_multi_speaker_simulator(config: DictConfig) -> tuple[CutSet, bool]: # Import here to avoid circular dependency from nemo.collections.asr.parts.utils.asr_multispeaker_utils import MultiSpeakerMixtureGenerator multi_speaker_cuts = CutSet( MultiSpeakerMixtureGenerator( manifest_filepath=config.manifest_filepath, simulator_type=config.simulator_type, sample_rate=config.get("sample_rate", 16000), min_delay=config.get("min_delay", 0.5), min_duration=config.get("min_duration", 0.1), max_duration=config.get("max_duration", 60), num_speakers=config.get("num_speakers", 2), global_rank=config.get("global_rank", 0), world_size=config.get("world_size", 1), ) ) is_tarred = config.get("is_tarred", False) return multi_speaker_cuts, is_tarred def mux( *cutsets: CutSet, weights: list[Union[int, float]], max_open_streams: Union[int, None] = None, seed: Union[str, int] = "trng", force_finite: bool = False, ) -> CutSet: """ Helper function to call the right multiplexing method flavour in lhotse. The result is always an infinitely iterable ``CutSet``, but depending on whether ``max_open_streams`` is set, it will select a more appropriate multiplexing strategy. """ if max_open_streams is not None: assert not force_finite, "max_open_streams and metadata_only/force_finite options are not compatible" cuts = CutSet.infinite_mux(*cutsets, weights=weights, seed=seed, max_open_streams=max_open_streams) else: if not force_finite: cutsets = [cs.repeat(preserve_id=True) for cs in cutsets] if len(cutsets) == 1: # CutSet.mux must take more than one CutSet. cuts = cutsets[0] else: cuts = CutSet.mux(*cutsets, weights=weights, seed=seed) return cuts def guess_parse_cutset(inp: Union[str, dict, omegaconf.DictConfig]) -> CutSet: """ Utility function that supports opening a CutSet from: * a string path to YAML input spec (see :func:`read_dataset_config` for details) * a string path to Lhotse non-tarred JSONL manifest * a string path to NeMo non-tarred JSON manifest * a dictionary specifying inputs with keys available in :class:`nemo.collections.common.data.lhotse.dataloader.LhotseDataLoadingConfig` It's intended to be used in a generic context where we are not sure which way the user will specify the inputs. """ from nemo.collections.common.data.lhotse.dataloader import make_structured_with_schema_warnings if isinstance(inp, (dict, omegaconf.DictConfig)): try: config = make_structured_with_schema_warnings(OmegaConf.from_dotlist([f"{k}={v}" for k, v in inp.items()])) cuts, _ = read_cutset_from_config(config) return cuts except Exception as e: raise RuntimeError( f"Couldn't open CutSet based on dict input {inp} (is it compatible with LhotseDataLoadingConfig?)" ) from e elif isinstance(inp, str): if inp.endswith(".yaml"): # Path to YAML file with the input configuration config = make_structured_with_schema_warnings(OmegaConf.from_dotlist([f"input_cfg={inp}"])) elif inp.endswith(".jsonl") or inp.endswith(".jsonl.gz"): # Path to a Lhotse non-tarred manifest config = make_structured_with_schema_warnings(OmegaConf.from_dotlist([f"cuts_path={inp}"])) else: # Assume anything else is a NeMo non-tarred manifest config = make_structured_with_schema_warnings(OmegaConf.from_dotlist([f"manifest_filepath={inp}"])) cuts, _ = read_cutset_from_config(config) return cuts else: raise RuntimeError(f'Unsupported input type: {type(inp)} (expected a dict or a string)') def _convert_tarred_to_duplex(cut, agent_silence_duration): """Helper function to convert single supervision to duplex format. This is a module-level function (not nested) so it can be pickled for multiprocessing. """ if len(cut.supervisions) != 1: # Skip cuts that don't have exactly one supervision return cut original_sup = cut.supervisions[0] orig_user_duration = cut.duration # Note here we use the last part of the audio as agent silence, which may cut user text, but this avoid using synthetic silence # TODO(kevinhu): Evaluate how this impacts user EOU # Append agent_silence_duration of silence to the original recording if agent_silence_duration > 0: sr = cut.recording.sampling_rate user_sil_samples = int(agent_silence_duration * sr) silence_audio = np.zeros((1, user_sil_samples), dtype=np.float32) # Concatenate silence to the end of the original audio orig_audio = cut.recording.load_audio() if orig_audio.ndim == 1: orig_audio = orig_audio[None, :] new_audio = np.concatenate([orig_audio, silence_audio], axis=1) # Create a new Recording with the extended audio new_recording = create_recording_from_array(new_audio, sr, cut.recording.id) cut.recording = new_recording cut.duration = new_audio.shape[1] / sr # Create user supervision (original speech) user_dur = orig_user_duration user_sup = SupervisionSegment( id=f"{cut.id}_user", recording_id=cut.recording_id, start=0.0, duration=user_dur, text=original_sup.text, language=original_sup.language, speaker="user", ) # Create agent supervision (silence with configurable duration) agent_start = orig_user_duration + agent_silence_duration if agent_silence_duration < 0 else orig_user_duration agent_dur = abs(agent_silence_duration) agent_sup = SupervisionSegment( id=f"{cut.id}_agent", recording_id=cut.recording_id, start=agent_start, duration=agent_dur, text="", # Empty text for silence language=original_sup.language, speaker="agent", ) # Create target_audio with all zeros (silence) sr = cut.recording.sampling_rate num_samples = int(cut.duration * sr) silence_audio = np.zeros((1, num_samples), dtype=np.float32) # Create a Recording from the silence audio silence_recording = create_recording_from_array(silence_audio, sr, f"{cut.id}_target") # Replace the single supervision with user and agent supervisions cut.supervisions = [user_sup, agent_sup] cut.task = "asr" # Add target_audio to cut.custom if cut.custom is None: cut.custom = {} cut.custom["target_audio"] = silence_recording return cut @data_type_parser(["nemo_tarred_to_duplex"]) def read_nemo_tarred_to_duplex(config) -> tuple[CutSet, bool]: """Convert single-supervision NeMo tarred data to duplex format with user speech and agent silence. This parser wraps ``nemo_tarred`` and converts each single-supervision cut into a two-supervision (user + agent) duplex cut. The user supervision keeps the original audio and text, while the agent supervision is silence with empty text. A silent ``target_audio`` recording (all zeros, same length as the source) is attached via ``cut.custom["target_audio"]``. Input config YAML example (used inside an ``input_cfg`` list):: - manifest_filepath: /path/to/manifest__OP_0..63_CL_.json tarred_audio_filepaths: /path/to/audio__OP_0..63_CL_.tar type: nemo_tarred_to_duplex agent_silence_duration: 2.0 # optional, default -0.08 weight: 1.0 Each line in the NeMo manifest JSON is the standard NeMo format:: {"audio_filepath": "relative/path.wav", "text": "transcript", "duration": 4.5} ``agent_silence_duration`` controls how much silence is used for the agent turn. A positive value appends that many seconds of silence after the user audio. A negative value (the default, ``-0.08``) re-uses the last ``abs(value)`` seconds of the user audio as the agent region instead of appending new silence. """ # by default, use the last part of user audio as agent silence duration agent_silence_duration = config.get("agent_silence_duration", -0.08) # Reuse the existing nemo_tarred parser by creating a config with type: nemo_tarred nemo_config = DictConfig(config) nemo_config.type = "nemo_tarred" # Load the cuts using the original parser cuts, is_tarred = read_nemo_manifest(nemo_config) # Apply the conversion using functools.partial to make it picklable convert_fn = partial(_convert_tarred_to_duplex, agent_silence_duration=agent_silence_duration) cuts = cuts.map(convert_fn) return cuts, is_tarred