wehub-resource-sync
ba4be087d5
CICD NeMo / cicd-main-unit-tests (push) Blocked by required conditions
CICD NeMo / cicd-main-speech (push) Blocked by required conditions
CICD NeMo / cicd-test-container-build (push) Blocked by required conditions
CICD NeMo / cicd-import-tests (push) Blocked by required conditions
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Blocked by required conditions
CICD NeMo / Nemo_CICD_Test (push) Blocked by required conditions
CICD NeMo / Coverage (e2e) (push) Blocked by required conditions
CICD NeMo / Coverage (unit-test) (push) Blocked by required conditions
CodeQL / Analyze (python) (push) Waiting to run
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
CICD NeMo / cicd-wait-in-queue (push) Waiting to run
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
3241 lines
114 KiB
Python
3241 lines
114 KiB
Python
# 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.
|
|
from collections import Counter
|
|
from io import BytesIO
|
|
from itertools import islice
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
import lhotse
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from lhotse import CutSet, MonoCut, NumpyFilesWriter, Recording, compute_num_samples
|
|
from lhotse.audio import AudioLoadingError
|
|
from lhotse.cut import Cut, MixedCut, PaddingCut
|
|
from lhotse.dataset import RoundRobinSampler, ZipSampler
|
|
from lhotse.shar import JsonlShardWriter
|
|
from lhotse.testing.dummies import dummy_recording
|
|
from lhotse.testing.random import deterministic_rng
|
|
from omegaconf import OmegaConf
|
|
|
|
from nemo.collections.common.data.lhotse import get_lhotse_dataloader_from_config
|
|
from nemo.collections.common.data.lhotse.text_adapters import SourceTargetTextExample, TextExample
|
|
from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer, create_spt_model
|
|
from nemo.utils.dependency import is_module_available
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def cutset_path(tmp_path_factory) -> Path:
|
|
"""10 utterances of length 1s as a Lhotse CutSet."""
|
|
from lhotse import CutSet
|
|
from lhotse.testing.dummies import DummyManifest
|
|
|
|
cuts = DummyManifest(CutSet, begin_id=0, end_id=10, with_data=True)
|
|
for c in cuts:
|
|
c.features = None
|
|
c.custom = None
|
|
c.supervisions[0].custom = None
|
|
|
|
tmp_path = tmp_path_factory.mktemp("data")
|
|
p = tmp_path / "cuts.jsonl.gz"
|
|
pa = tmp_path / "audio"
|
|
cuts.save_audios(pa).to_file(p)
|
|
return p
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def cutset_shar_path(cutset_path: Path) -> Path:
|
|
"""10 utterances of length 1s as a Lhotse Shar (tarred) CutSet."""
|
|
from lhotse import CutSet
|
|
|
|
cuts = CutSet.from_file(cutset_path)
|
|
p = cutset_path.parent / "shar"
|
|
p.mkdir(exist_ok=True)
|
|
cuts.to_shar(p, fields={"recording": "wav"}, shard_size=5)
|
|
return p
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def cutset_shar_path_other(cutset_path: Path) -> Path:
|
|
"""10 utterances of length 1s as a Lhotse Shar (tarred) CutSet, but with different IDs."""
|
|
from lhotse import CutSet
|
|
|
|
cuts = CutSet.from_file(cutset_path).modify_ids(lambda id: f"other-{id}")
|
|
p = cutset_path.parent / "shar-other"
|
|
p.mkdir(exist_ok=True)
|
|
cuts.to_shar(p, fields={"recording": "wav"}, shard_size=5)
|
|
return p
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def nemo_manifest_path(cutset_path: Path):
|
|
"""10 utterances of length 1s as a NeMo manifest."""
|
|
from lhotse import CutSet
|
|
from lhotse.serialization import save_to_jsonl
|
|
|
|
nemo = []
|
|
for c in CutSet.from_file(cutset_path):
|
|
nemo.append(
|
|
{
|
|
"audio_filepath": c.recording.sources[0].source,
|
|
"text": "irrelevant",
|
|
"text-other": "not relevant",
|
|
"duration": c.duration,
|
|
"my-custom-field": "irrelevant",
|
|
"lang": "en",
|
|
"custom-lang": "pl",
|
|
}
|
|
)
|
|
p = cutset_path.parent / "nemo_manifest.json"
|
|
save_to_jsonl(nemo, p)
|
|
return p
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def nemo_manifest_with_skipme_path(nemo_manifest_path: Path) -> Path:
|
|
"""Create a nemo manifest with last 2 utterances out of 10 with `_skipme` key enabled"""
|
|
from lhotse.serialization import load_jsonl, save_to_jsonl
|
|
|
|
all_items = list(load_jsonl(nemo_manifest_path))
|
|
|
|
for item in all_items[-2:]:
|
|
item['_skipme'] = True
|
|
|
|
p = nemo_manifest_path.parent / "nemo_manifest_with_skipme.json"
|
|
save_to_jsonl(all_items, p)
|
|
return p
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def mc_cutset_path(tmp_path_factory) -> Path:
|
|
"""10 two-channel utterances of length 1s as a Lhotse CutSet."""
|
|
from lhotse import CutSet, MultiCut
|
|
from lhotse.testing.dummies import DummyManifest
|
|
|
|
num_examples = 10 # number of examples
|
|
num_channels = 2 # number of channels per example
|
|
|
|
# create a dummy manifest with single-channel examples
|
|
sc_cuts = DummyManifest(CutSet, begin_id=0, end_id=num_examples * num_channels, with_data=True)
|
|
mc_cuts = []
|
|
|
|
for n in range(num_examples):
|
|
# sources for individual channels
|
|
mc_sources = []
|
|
for channel in range(num_channels):
|
|
source = sc_cuts[n * num_channels + channel].recording.sources[0]
|
|
source.channels = [channel]
|
|
mc_sources.append(source)
|
|
|
|
# merge recordings
|
|
rec = Recording(
|
|
sources=mc_sources,
|
|
id=f'mc-dummy-recording-{n:02d}',
|
|
num_samples=sc_cuts[0].num_samples,
|
|
duration=sc_cuts[0].duration,
|
|
sampling_rate=sc_cuts[0].sampling_rate,
|
|
)
|
|
|
|
# multi-channel cut
|
|
cut = MultiCut(
|
|
recording=rec, id=f'mc-dummy-cut-{n:02d}', start=0, duration=1.0, channel=list(range(num_channels))
|
|
)
|
|
mc_cuts.append(cut)
|
|
|
|
mc_cuts = CutSet.from_cuts(mc_cuts)
|
|
|
|
tmp_path = tmp_path_factory.mktemp("data")
|
|
p = tmp_path / "mc_cuts.jsonl.gz"
|
|
pa = tmp_path / "mc_audio"
|
|
mc_cuts.save_audios(pa).to_file(p)
|
|
return p
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def nemo_tarred_manifest_path(nemo_manifest_path: Path) -> Tuple[str, str]:
|
|
"""10 utterances of length 1s as a NeMo tarred manifest."""
|
|
from lhotse.serialization import SequentialJsonlWriter, load_jsonl
|
|
from lhotse.shar.writers import TarWriter
|
|
|
|
root = nemo_manifest_path.parent / "nemo_tar"
|
|
root.mkdir(exist_ok=True)
|
|
|
|
with (
|
|
TarWriter(f"{root}/audios_%01d.tar", shard_size=5) as tar_writer,
|
|
SequentialJsonlWriter(root / "tarred_audio_filepaths.jsonl") as mft_writer,
|
|
):
|
|
for idx, d in enumerate(load_jsonl(nemo_manifest_path)):
|
|
p = d["audio_filepath"]
|
|
name = Path(p).name
|
|
with open(p, "rb") as f:
|
|
tar_writer.write(name, BytesIO(f.read()))
|
|
mft_writer.write({**d, "audio_filepath": name, "shard_id": int(idx > 4)})
|
|
return mft_writer.path, f"{root}/audios__OP_0..1_CL_.tar"
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def nemo_tarred_manifest_with_skipme_path(nemo_tarred_manifest_path: Path) -> Tuple[str, str]:
|
|
"""Create a nemo tarred manifest with last 2 utterances out of 10 with `_skipme` key enabled."""
|
|
from lhotse.serialization import load_jsonl, save_to_jsonl
|
|
|
|
json_p, tar_p = nemo_tarred_manifest_path
|
|
|
|
all_items = list(load_jsonl(json_p))
|
|
|
|
for item in all_items[-2:]:
|
|
item['_skipme'] = True
|
|
|
|
p = json_p.parent / "tarred_audio_filepaths_with_skipme.jsonl"
|
|
save_to_jsonl(all_items, p)
|
|
|
|
return p, tar_p
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def nemo_tarred_manifest_path_multi(nemo_tarred_manifest_path: tuple[str, str]) -> Tuple[str, str]:
|
|
"""10 utterances of length 1s as a NeMo tarred manifest. Stored in one manifest per shard."""
|
|
from lhotse.serialization import load_jsonl
|
|
from lhotse.shar.writers import JsonlShardWriter
|
|
|
|
json_p, tar_p = nemo_tarred_manifest_path
|
|
|
|
json_dir = json_p.parent / "shard_manifests"
|
|
json_dir.mkdir(exist_ok=True)
|
|
with JsonlShardWriter(f"{json_dir}/manifest_%d.jsonl", shard_size=5) as mft_writer:
|
|
for item in load_jsonl(json_p):
|
|
mft_writer.write(item)
|
|
return f"{json_dir}/manifest__OP_0..1_CL_.jsonl", tar_p
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def nemo_tarred_manifest_subset_path(nemo_tarred_manifest_path: Tuple[str, str]) -> Tuple[str, str]:
|
|
"""Create a shard manifests with randomly chosen 50% percent of tarred contents."""
|
|
from lhotse.serialization import load_jsonl
|
|
from lhotse.shar.writers import JsonlShardWriter
|
|
|
|
json_p, tar_p = nemo_tarred_manifest_path
|
|
json_dir = json_p.parent / "shard_manifests_subset"
|
|
json_dir.mkdir(exist_ok=True)
|
|
all_items = list(load_jsonl(json_p))
|
|
tarr_0_data = all_items[:5]
|
|
tarr_1_data = all_items[5:]
|
|
|
|
subset_items = tarr_0_data[-3:] + tarr_1_data[-3:]
|
|
with JsonlShardWriter(f"{json_dir}/manifest_%d.jsonl", shard_size=3) as mft_writer:
|
|
for item in subset_items:
|
|
mft_writer.write(item)
|
|
return f"{json_dir}/manifest__OP_0..1_CL_.jsonl", tar_p, subset_items
|
|
|
|
|
|
class UnsupervisedAudioDataset(torch.utils.data.Dataset):
|
|
def __getitem__(self, cuts: lhotse.CutSet) -> Dict[str, torch.Tensor]:
|
|
audio, audio_lens = lhotse.dataset.collation.collate_audio(cuts)
|
|
return {"audio": audio, "audio_lens": audio_lens, "ids": [c.id for c in cuts]}
|
|
|
|
|
|
def test_dataloader_from_lhotse_cuts(cutset_path: Path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": cutset_path,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
# lhotse specific
|
|
"use_bucketing": True,
|
|
"concurrent_bucketing": False,
|
|
"num_buckets": 2,
|
|
"drop_last": False,
|
|
"batch_duration": 4.0, # seconds
|
|
"quadratic_duration": 15.0, # seconds
|
|
"shuffle_buffer_size": 10,
|
|
"bucket_buffer_size": 100,
|
|
"seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
batches = [batch for batch in dl]
|
|
assert len(batches) == 4
|
|
|
|
b = batches[0]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[1]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[2]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[3]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 1
|
|
|
|
|
|
def test_dataloader_from_lhotse_cuts_truncate(cutset_path: Path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": cutset_path,
|
|
"truncate_duration": 0.5,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
batches = [b for b in dl]
|
|
assert len(batches) == 3
|
|
# 0.5s = 8000 samples, note the constant duration and batch size except for last batch
|
|
assert batches[0]["audio"].shape == (4, 8000)
|
|
assert batches[1]["audio"].shape == (4, 8000)
|
|
assert batches[2]["audio"].shape == (2, 8000)
|
|
# exactly 10 cuts were used
|
|
|
|
|
|
def test_dataloader_from_lhotse_cuts_cut_into_windows(cutset_path: Path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": cutset_path,
|
|
"cut_into_windows_duration": 0.5,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
batches = [b for b in dl]
|
|
assert len(batches) == 5
|
|
# 0.5s = 8000 samples, note the constant duration and batch size
|
|
assert batches[0]["audio"].shape == (4, 8000)
|
|
assert batches[1]["audio"].shape == (4, 8000)
|
|
assert batches[2]["audio"].shape == (4, 8000)
|
|
assert batches[3]["audio"].shape == (4, 8000)
|
|
assert batches[4]["audio"].shape == (4, 8000)
|
|
# exactly 20 cuts were used because we cut 10x 1s cuts into 20x 0.5s cuts
|
|
|
|
|
|
def test_dataloader_from_lhotse_cuts_pad_min_duration(cutset_path: Path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": cutset_path,
|
|
"pad_min_duration": 21.0,
|
|
"pad_direction": "left",
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 1,
|
|
"seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
batch = next(iter(dl))
|
|
(cut,) = batch
|
|
assert cut.duration == 21.0
|
|
assert isinstance(cut, MixedCut)
|
|
assert len(cut.tracks) == 2
|
|
assert isinstance(cut.tracks[0].cut, PaddingCut)
|
|
assert isinstance(cut.tracks[1].cut, MonoCut)
|
|
|
|
|
|
def test_dataloader_from_lhotse_cuts_channel_selector(mc_cutset_path: Path):
|
|
# Dataloader without channel selector
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": mc_cutset_path,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
batches = [b for b in dl]
|
|
assert len(batches) == 3
|
|
|
|
# 1.0s = 16000 samples, two channels, note the constant duration and batch size
|
|
assert batches[0]["audio"].shape == (4, 2, 16000)
|
|
assert batches[1]["audio"].shape == (4, 2, 16000)
|
|
assert batches[2]["audio"].shape == (2, 2, 16000)
|
|
# exactly 10 cuts were used
|
|
|
|
# Apply channel selector
|
|
for channel_selector in [None, 0, 1]:
|
|
|
|
config_cs = OmegaConf.create(
|
|
{
|
|
"cuts_path": mc_cutset_path,
|
|
"channel_selector": channel_selector,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
}
|
|
)
|
|
|
|
dl_cs = get_lhotse_dataloader_from_config(
|
|
config=config_cs, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
for n, b_cs in enumerate(dl_cs):
|
|
if channel_selector is None:
|
|
# no channel selector, needs to match the original dataset
|
|
assert torch.equal(b_cs["audio"], batches[n]["audio"])
|
|
else:
|
|
# channel selector, needs to match the selected channel
|
|
assert torch.equal(b_cs["audio"], batches[n]["audio"][:, channel_selector, :])
|
|
|
|
|
|
def test_dataloader_from_lhotse_shar_cuts(cutset_shar_path: Path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"shar_path": cutset_shar_path,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
# lhotse specific
|
|
"use_bucketing": True,
|
|
"concurrent_bucketing": False,
|
|
"num_buckets": 2,
|
|
"drop_last": False,
|
|
"batch_duration": 4.0, # seconds
|
|
"quadratic_duration": 15.0, # seconds
|
|
"shuffle_buffer_size": 10,
|
|
"bucket_buffer_size": 100,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
# Note: we use islice here because with Lhotse Shar the dataloader will always be infinite.
|
|
batches = [batch for batch in islice(dl, 4)]
|
|
assert len(batches) == 4
|
|
|
|
b = batches[0]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[1]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[2]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[3]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
|
|
def test_dataloader_from_lhotse_shar_cuts_via_fields(cutset_shar_path: Path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"shar_path": {
|
|
"cuts": f"{cutset_shar_path}/cuts._OP_000000..000001_CL_.jsonl.gz",
|
|
"recording": f"{cutset_shar_path}/recording._OP_000000..000001_CL_.tar",
|
|
},
|
|
"sample_rate": 16000,
|
|
"num_workers": 0,
|
|
"shuffle": False,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
batch = next(iter(dl))
|
|
assert len(batch) == 4
|
|
audio = batch[0].load_audio()
|
|
assert isinstance(audio, np.ndarray)
|
|
|
|
|
|
def test_dataloader_from_lhotse_shar_cuts_add_new_field(tmp_path_factory, cutset_shar_path: Path):
|
|
|
|
# We're creating a new field called "wer" that will be dynamically attached to Lhotse Shar cuts.
|
|
# Each "wer" shard is a jsonl manifest that has to match the "cuts" sharded manifest.
|
|
# It must have a "cut_id" field used for runtime check that the user provided correct paths.
|
|
# "wer" will be attached to each cut under `cut.wer` / cut.custom["wer"].
|
|
wer_dir = tmp_path_factory.mktemp("wer_dir")
|
|
with JsonlShardWriter(f"{wer_dir}/wer.%06d.jsonl.gz", shard_size=5) as writer:
|
|
for i in range(10):
|
|
writer.write({"cut_id": "dummy-mono-cut-%04d" % i, "wer": 0.5})
|
|
|
|
config = OmegaConf.create(
|
|
{
|
|
"shar_path": {
|
|
"cuts": f"{cutset_shar_path}/cuts._OP_000000..000001_CL_.jsonl.gz",
|
|
"recording": f"{cutset_shar_path}/recording._OP_000000..000001_CL_.tar",
|
|
"wer": f"{wer_dir}/wer._OP_000000..000001_CL_.jsonl.gz",
|
|
},
|
|
"sample_rate": 16000,
|
|
"num_workers": 0,
|
|
"shuffle": False,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
batch = next(iter(dl))
|
|
assert len(batch) == 4
|
|
assert batch[0].wer == 0.5
|
|
|
|
|
|
def test_dataloader_from_nemo_manifest(nemo_manifest_path: Path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": nemo_manifest_path,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
# lhotse specific
|
|
"use_bucketing": True,
|
|
"concurrent_bucketing": False,
|
|
"num_buckets": 2,
|
|
"drop_last": False,
|
|
"batch_duration": 4.0, # seconds
|
|
"quadratic_duration": 15.0, # seconds
|
|
"shuffle_buffer_size": 10,
|
|
"bucket_buffer_size": 100,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
batches = [batch for batch in dl]
|
|
assert len(batches) == 4
|
|
|
|
b = batches[0]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[1]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[2]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[3]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 1
|
|
|
|
|
|
class _Identity:
|
|
def __getitem__(self, cuts):
|
|
return cuts
|
|
|
|
|
|
def test_dataloader_from_nemo_manifest_has_custom_fields(nemo_manifest_path: Path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": nemo_manifest_path,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
# lhotse specific
|
|
"use_bucketing": False,
|
|
"batch_duration": 4.0, # seconds
|
|
"shuffle_buffer_size": 10,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=_Identity())
|
|
|
|
batch = next(iter(dl))
|
|
for cut in batch:
|
|
assert isinstance(cut.custom, dict)
|
|
assert "my-custom-field" in cut.custom
|
|
|
|
|
|
def test_dataloader_from_tarred_nemo_manifest(nemo_tarred_manifest_path: tuple[str, str]):
|
|
json_mft, tar_mft = nemo_tarred_manifest_path
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": json_mft,
|
|
"tarred_audio_filepaths": tar_mft,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
# lhotse specific
|
|
"use_bucketing": True,
|
|
"concurrent_bucketing": False,
|
|
"num_buckets": 2,
|
|
"drop_last": False,
|
|
"batch_duration": 4.0, # seconds
|
|
"quadratic_duration": 15.0, # seconds
|
|
"shuffle_buffer_size": 10,
|
|
"bucket_buffer_size": 100,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
batches = [batch for batch in islice(dl, 4)]
|
|
assert len(batches) == 4
|
|
|
|
b = batches[0]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[1]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[2]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[3]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
|
|
def test_dataloader_from_tarred_nemo_manifest_weighted_combination(nemo_tarred_manifest_path: tuple[str, str]):
|
|
json_mft, tar_mft = nemo_tarred_manifest_path
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": [[json_mft, 0.8], [json_mft, 0.2]],
|
|
"tarred_audio_filepaths": [[tar_mft], [tar_mft]],
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
# lhotse specific
|
|
"use_bucketing": True,
|
|
"concurrent_bucketing": False,
|
|
"num_buckets": 2,
|
|
"drop_last": False,
|
|
"batch_duration": 4.0, # seconds
|
|
"quadratic_duration": 15.0, # seconds
|
|
"shuffle_buffer_size": 10,
|
|
"bucket_buffer_size": 100,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
b = next(iter(dl))
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
|
|
def test_dataloader_from_tarred_nemo_manifest_multi(nemo_tarred_manifest_path_multi: tuple[str, str]):
|
|
json_mft, tar_mft = nemo_tarred_manifest_path_multi
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": json_mft,
|
|
"tarred_audio_filepaths": tar_mft,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
# lhotse specific
|
|
"use_bucketing": True,
|
|
"concurrent_bucketing": False,
|
|
"num_buckets": 2,
|
|
"drop_last": False,
|
|
"batch_duration": 4.0, # seconds
|
|
"quadratic_duration": 15.0, # seconds
|
|
"shuffle_buffer_size": 10,
|
|
"bucket_buffer_size": 100,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
batches = [batch for batch in islice(dl, 4)]
|
|
assert len(batches) == 4
|
|
|
|
b = batches[0]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[1]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[2]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
b = batches[3]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 3
|
|
|
|
|
|
def test_dataloader_from_tarred_nemo_manifest_multi_max_open_streams(nemo_tarred_manifest_path_multi: tuple[str, str]):
|
|
json_mft, tar_mft = nemo_tarred_manifest_path_multi
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": [[json_mft], [json_mft]],
|
|
"tarred_audio_filepaths": [[tar_mft], [tar_mft]],
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
# lhotse specific
|
|
"use_bucketing": True,
|
|
"concurrent_bucketing": False,
|
|
"num_buckets": 2,
|
|
"max_open_streams": 1,
|
|
"drop_last": False,
|
|
"batch_duration": 4.0, # seconds
|
|
"quadratic_duration": 15.0, # seconds
|
|
"shuffle_buffer_size": 10,
|
|
"bucket_buffer_size": 100,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
_ = next(iter(dl))
|
|
|
|
|
|
@pytest.mark.skipif(not is_module_available("pyloudnorm"), reason="pyloudnorm is required to concatenate samples")
|
|
def test_dataloader_from_tarred_nemo_manifest_concat(nemo_tarred_manifest_path: tuple[str, str]):
|
|
json_mft, tar_mft = nemo_tarred_manifest_path
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": json_mft,
|
|
"tarred_audio_filepaths": tar_mft,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
# lhotse specific
|
|
"concatenate_samples": True,
|
|
"concatenate_duration_factor": 3.0,
|
|
"batch_duration": 4.0,
|
|
"quadratic_duration": 15.0, # seconds
|
|
"use_bucketing": False,
|
|
"drop_last": False,
|
|
"shuffle_buffer_size": 10,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
batches = [batch for batch in islice(dl, 4)]
|
|
|
|
assert len(batches) == 4
|
|
|
|
# the first element has been concatenated: 2x16000 speech (2x1s) + 1600 gap (0.1s)
|
|
expected_audio_lens = torch.tensor([33600, 16000], dtype=torch.int32)
|
|
|
|
b = batches[0]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 2
|
|
torch.testing.assert_close(b["audio_lens"], expected_audio_lens)
|
|
|
|
b = batches[1]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 2
|
|
torch.testing.assert_close(b["audio_lens"], expected_audio_lens)
|
|
|
|
b = batches[2]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 2
|
|
torch.testing.assert_close(b["audio_lens"], expected_audio_lens)
|
|
|
|
b = batches[3]
|
|
assert set(b.keys()) == {"audio", "audio_lens", "ids"}
|
|
assert b["audio"].shape[0] == b["audio_lens"].shape[0] == 2
|
|
torch.testing.assert_close(b["audio_lens"], expected_audio_lens)
|
|
|
|
|
|
def test_dataloader_from_lhotse_shar_cuts_combine_datasets_unweighted(
|
|
cutset_shar_path: Path, cutset_shar_path_other: Path
|
|
):
|
|
"""
|
|
Note: if we iterated more mini-batches in this test, in the expectation there
|
|
will be 50-50 % mini-batch occupancy of examples from both datasets.
|
|
"""
|
|
config = OmegaConf.create(
|
|
{
|
|
"shar_path": [cutset_shar_path, cutset_shar_path_other],
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
# lhotse specific
|
|
"use_bucketing": True,
|
|
"concurrent_bucketing": False,
|
|
"num_buckets": 2,
|
|
"drop_last": False,
|
|
"batch_duration": 4.0, # seconds
|
|
"quadratic_duration": 15.0, # seconds
|
|
"shuffle_buffer_size": 10,
|
|
"bucket_buffer_size": 100,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
# Note: we use islice here because with Lhotse Shar the dataloader will always be infinite.
|
|
batches = [batch for batch in islice(dl, 4)]
|
|
assert len(batches) == 4
|
|
|
|
b = batches[0]
|
|
assert len([cid for cid in b["ids"] if cid.startswith("dummy")]) == 1 # dataset 1
|
|
assert len([cid for cid in b["ids"] if cid.startswith("other")]) == 2 # dataset 2
|
|
|
|
b = batches[1]
|
|
assert len([cid for cid in b["ids"] if cid.startswith("dummy")]) == 0 # dataset 1
|
|
assert len([cid for cid in b["ids"] if cid.startswith("other")]) == 3 # dataset 2
|
|
|
|
b = batches[2]
|
|
assert len([cid for cid in b["ids"] if cid.startswith("dummy")]) == 2 # dataset 1
|
|
assert len([cid for cid in b["ids"] if cid.startswith("other")]) == 1 # dataset 2
|
|
|
|
b = batches[3]
|
|
assert len([cid for cid in b["ids"] if cid.startswith("dummy")]) == 1 # dataset 1
|
|
assert len([cid for cid in b["ids"] if cid.startswith("other")]) == 2 # dataset 2
|
|
|
|
|
|
def test_dataloader_from_lhotse_shar_cuts_combine_datasets_weighted(
|
|
cutset_shar_path: Path, cutset_shar_path_other: Path
|
|
):
|
|
"""
|
|
Note: if we iterated more mini-batches in this test, in the expectation there
|
|
will be 90-10 % mini-batch occupancy of examples from both datasets.
|
|
"""
|
|
config = OmegaConf.create(
|
|
{
|
|
"shar_path": [[cutset_shar_path, 90], [cutset_shar_path_other, 10]],
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
# lhotse specific
|
|
"use_bucketing": True,
|
|
"concurrent_bucketing": False,
|
|
"num_buckets": 2,
|
|
"drop_last": False,
|
|
"batch_duration": 4.0, # seconds
|
|
"quadratic_duration": 15.0, # seconds
|
|
"shuffle_buffer_size": 10,
|
|
"bucket_buffer_size": 100,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
# Note: we use islice here because with Lhotse Shar the dataloader will always be infinite.
|
|
batches = [batch for batch in islice(dl, 6)]
|
|
assert len(batches) == 6
|
|
|
|
b = batches[0]
|
|
assert len([cid for cid in b["ids"] if cid.startswith("dummy")]) == 3 # dataset 1
|
|
assert len([cid for cid in b["ids"] if cid.startswith("other")]) == 0 # dataset 2
|
|
|
|
b = batches[1]
|
|
assert len([cid for cid in b["ids"] if cid.startswith("dummy")]) == 1 # dataset 1
|
|
assert len([cid for cid in b["ids"] if cid.startswith("other")]) == 2 # dataset 2
|
|
|
|
b = batches[2]
|
|
assert len([cid for cid in b["ids"] if cid.startswith("dummy")]) == 2 # dataset 1
|
|
assert len([cid for cid in b["ids"] if cid.startswith("other")]) == 1 # dataset 2
|
|
|
|
b = batches[3]
|
|
assert len([cid for cid in b["ids"] if cid.startswith("dummy")]) == 3 # dataset 1
|
|
assert len([cid for cid in b["ids"] if cid.startswith("other")]) == 0 # dataset 2
|
|
|
|
b = batches[4]
|
|
assert len([cid for cid in b["ids"] if cid.startswith("dummy")]) == 3 # dataset 1
|
|
assert len([cid for cid in b["ids"] if cid.startswith("other")]) == 0 # dataset 2
|
|
|
|
b = batches[5]
|
|
assert len([cid for cid in b["ids"] if cid.startswith("dummy")]) == 3 # dataset 1
|
|
assert len([cid for cid in b["ids"] if cid.startswith("other")]) == 0 # dataset 2
|
|
|
|
|
|
class TextDataset(torch.utils.data.Dataset):
|
|
def __getitem__(self, cuts: lhotse.CutSet) -> List[str]:
|
|
return [c.supervisions[0].text for c in cuts]
|
|
|
|
|
|
@pytest.mark.parametrize(["text_field", "text_value"], [(None, "irrelevant"), ("text-other", "not relevant")])
|
|
def test_dataloader_from_nemo_manifest_with_text_field(nemo_manifest_path: Path, text_field: str, text_value: str):
|
|
kwarg = {"text_field": text_field} if text_field is not None else {}
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": nemo_manifest_path,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 2,
|
|
# lhotse specific
|
|
"use_bucketing": False,
|
|
**kwarg,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=TextDataset())
|
|
b = next(iter(dl))
|
|
assert b == [text_value] * 2
|
|
|
|
|
|
class LangDataset(torch.utils.data.Dataset):
|
|
def __getitem__(self, cuts: lhotse.CutSet) -> List[str]:
|
|
return [c.supervisions[0].language for c in cuts]
|
|
|
|
|
|
@pytest.mark.parametrize(["lang_field", "lang_value"], [(None, "en"), ("custom-lang", "pl")])
|
|
def test_dataloader_from_nemo_manifest_with_lang_field(nemo_manifest_path: Path, lang_field: str, lang_value: str):
|
|
kwarg = {"lang_field": lang_field} if lang_field is not None else {}
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": nemo_manifest_path,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 2,
|
|
# lhotse specific
|
|
"use_bucketing": False,
|
|
**kwarg,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=LangDataset())
|
|
b = next(iter(dl))
|
|
assert b == [lang_value] * 2
|
|
|
|
|
|
def test_lazy_nemo_iterator_with_offset_field(tmp_path: Path):
|
|
import numpy as np
|
|
import soundfile as sf
|
|
|
|
from nemo.collections.common.data.lhotse.nemo_adapters import LazyNeMoIterator
|
|
|
|
# Have to generate as INT16 to avoid quantization error after saving to 16-bit WAV
|
|
INT16MAX = 2**15
|
|
expected_audio = np.random.randint(low=-INT16MAX - 1, high=INT16MAX, size=(16000,)).astype(np.float32) / INT16MAX
|
|
audio_path = str(tmp_path / "dummy.wav")
|
|
sf.write(audio_path, expected_audio, 16000)
|
|
|
|
manifest_path = str(tmp_path / "manifest.json")
|
|
lhotse.serialization.save_to_jsonl(
|
|
[
|
|
{"audio_filepath": audio_path, "offset": 0.0, "duration": 0.5, "text": "irrelevant"},
|
|
{"audio_filepath": audio_path, "offset": 0.5, "duration": 0.5, "text": "irrelevant"},
|
|
],
|
|
manifest_path,
|
|
)
|
|
|
|
cuts = lhotse.CutSet(LazyNeMoIterator(manifest_path))
|
|
|
|
cut = cuts[0]
|
|
assert isinstance(cut, lhotse.MonoCut)
|
|
assert cut.start == 0.0
|
|
assert cut.duration == 0.5
|
|
assert cut.sampling_rate == 16000
|
|
assert cut.num_samples == 8000
|
|
assert cut.supervisions[0].text == "irrelevant"
|
|
audio = cut.load_audio()
|
|
assert audio.shape == (1, 8000)
|
|
np.testing.assert_equal(audio[0], expected_audio[:8000])
|
|
|
|
cut = cuts[1]
|
|
assert isinstance(cut, lhotse.MonoCut)
|
|
assert cut.start == 0.5
|
|
assert cut.duration == 0.5
|
|
assert cut.sampling_rate == 16000
|
|
assert cut.num_samples == 8000
|
|
assert cut.supervisions[0].text == "irrelevant"
|
|
audio = cut.load_audio()
|
|
assert audio.shape == (1, 8000)
|
|
np.testing.assert_allclose(audio[0], expected_audio[8000:], atol=5e-5)
|
|
|
|
assert cuts[0].id != cuts[1].id
|
|
|
|
|
|
def test_lazy_nemo_iterator_with_relative_paths(tmp_path: Path):
|
|
import numpy as np
|
|
import soundfile as sf
|
|
|
|
from nemo.collections.common.data.lhotse.nemo_adapters import LazyNeMoIterator
|
|
|
|
# Have to generate as INT16 to avoid quantization error after saving to 16-bit WAV
|
|
INT16MAX = 2**15
|
|
expected_audio = np.random.randint(low=-INT16MAX - 1, high=INT16MAX, size=(16000,)).astype(np.float32) / INT16MAX
|
|
audio_path = str(tmp_path / "dummy.wav")
|
|
sf.write(audio_path, expected_audio, 16000)
|
|
|
|
manifest_path = str(tmp_path / "manifest.json")
|
|
lhotse.serialization.save_to_jsonl(
|
|
[
|
|
# note: relative path
|
|
{"audio_filepath": "dummy.wav", "offset": 0.0, "duration": 0.5, "text": "irrelevant"},
|
|
],
|
|
manifest_path,
|
|
)
|
|
|
|
cuts = lhotse.CutSet(LazyNeMoIterator(manifest_path))
|
|
cut = cuts[0]
|
|
audio = cut.load_audio()
|
|
|
|
assert isinstance(cut, lhotse.MonoCut)
|
|
assert cut.start == 0.0
|
|
assert cut.duration == 0.5
|
|
assert cut.sampling_rate == 16000
|
|
assert cut.num_samples == 8000
|
|
assert cut.supervisions[0].text == "irrelevant"
|
|
assert audio.shape == (1, 8000)
|
|
np.testing.assert_equal(audio[0], expected_audio[:8000])
|
|
|
|
|
|
def test_lhotse_cuts_resolve_relative_paths(tmp_path: Path):
|
|
cuts_path = tmp_path / "cuts.jsonl.gz"
|
|
audio_path = tmp_path / "_relative_test_audio_.wav"
|
|
lhotse.audio.save_audio(audio_path, np.random.rand(16000) - 0.5, 16000)
|
|
cut = Recording.from_file(audio_path).to_cut()
|
|
cut.recording.sources[0].source = str(audio_path.name) # make the path relative
|
|
cut.target_recording = cut.recording # assign a custom field with relative path
|
|
with NumpyFilesWriter(tmp_path) as w:
|
|
cut.some_array = w.store_array(cut.id, np.random.randn(32))
|
|
cut.some_array.storage_path = "" # relative path
|
|
|
|
with pytest.raises(AudioLoadingError):
|
|
cut.load_audio() # Lhotse doesn't know about what the path should be relative to
|
|
cut.load_target_recording()
|
|
|
|
CutSet([cut]).to_file(cuts_path)
|
|
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": cuts_path,
|
|
"sample_rate": 16000,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 2,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=_Identity())
|
|
|
|
batches = [batch for batch in dl]
|
|
assert len(batches) == 1
|
|
|
|
for cut in batches[0]:
|
|
assert cut.has_recording
|
|
cut.load_audio() # works
|
|
assert cut.has_custom("target_recording")
|
|
cut.load_target_recording()
|
|
assert cut.has_custom("some_array")
|
|
cut.load_some_array()
|
|
|
|
|
|
class Identity(torch.utils.data.Dataset):
|
|
def __getitem__(self, cuts: lhotse.CutSet) -> lhotse.CutSet:
|
|
return cuts
|
|
|
|
|
|
def test_extended_data_input_cfg(cutset_shar_path, nemo_tarred_manifest_path_multi):
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
"type": "nemo_tarred",
|
|
"manifest_filepath": nemo_tarred_manifest_path_multi[0],
|
|
"tarred_audio_filepaths": nemo_tarred_manifest_path_multi[1],
|
|
"weight": 0.5,
|
|
"tags": {
|
|
"language": "en",
|
|
"modality": "audio",
|
|
"dataset_name": "D1",
|
|
},
|
|
},
|
|
{
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path,
|
|
"weight": 0.5,
|
|
"tags": {
|
|
"language": "en",
|
|
"modality": "audio",
|
|
"dataset_name": "D2",
|
|
},
|
|
},
|
|
],
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"num_workers": 0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
# Note: we use islice here because the dataloader will be infinite.
|
|
batches = [batch for batch in islice(dl, 2)]
|
|
|
|
b = batches[0]
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert all(c.custom["language"] == "en" for c in b)
|
|
assert all(c.custom["modality"] == "audio" for c in b)
|
|
assert sum(c.custom["dataset_name"] == "D1" for c in b) == 2
|
|
assert sum(c.custom["dataset_name"] == "D2" for c in b) == 2
|
|
|
|
b = batches[1]
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert all(c.custom["language"] == "en" for c in b)
|
|
assert all(c.custom["modality"] == "audio" for c in b)
|
|
assert sum(c.custom["dataset_name"] == "D1" for c in b) == 1
|
|
assert sum(c.custom["dataset_name"] == "D2" for c in b) == 3
|
|
|
|
|
|
def test_extended_data_input_cfg_subgroup(cutset_shar_path, nemo_tarred_manifest_path_multi):
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
"type": "group",
|
|
"input_cfg": [
|
|
{
|
|
"type": "nemo_tarred",
|
|
"manifest_filepath": nemo_tarred_manifest_path_multi[0],
|
|
"tarred_audio_filepaths": nemo_tarred_manifest_path_multi[1],
|
|
"weight": 0.5,
|
|
"tags": {
|
|
"language": "en",
|
|
"modality": "audio",
|
|
"dataset_name": "D1",
|
|
},
|
|
},
|
|
{
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path,
|
|
"weight": 0.5,
|
|
"tags": {
|
|
"language": "en",
|
|
"modality": "audio",
|
|
"dataset_name": "D2",
|
|
},
|
|
},
|
|
],
|
|
"weight": 0.2,
|
|
"tags": {
|
|
"group_name": "G1",
|
|
},
|
|
},
|
|
{
|
|
"type": "group",
|
|
"weight": 0.8,
|
|
"input_cfg": [
|
|
{
|
|
"type": "nemo_tarred",
|
|
"manifest_filepath": nemo_tarred_manifest_path_multi[0],
|
|
"tarred_audio_filepaths": nemo_tarred_manifest_path_multi[1],
|
|
"weight": 0.5,
|
|
"tags": {
|
|
"language": "en",
|
|
"modality": "audio",
|
|
"dataset_name": "D3",
|
|
},
|
|
},
|
|
{
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path,
|
|
"weight": 0.5,
|
|
"tags": {
|
|
"language": "en",
|
|
"modality": "audio",
|
|
"dataset_name": "D4",
|
|
},
|
|
},
|
|
],
|
|
"tags": {
|
|
"group_name": "G2",
|
|
},
|
|
},
|
|
],
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"num_workers": 0,
|
|
"batch_size": 32,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
# Sample 100 mini-batches and test statistical properties
|
|
group_occurrences = Counter()
|
|
dataset_occurrences = Counter()
|
|
for batch in islice(dl, 100):
|
|
for cut in batch:
|
|
group_occurrences[cut.group_name] += 1
|
|
dataset_occurrences[cut.dataset_name] += 1
|
|
|
|
tot = sum(group_occurrences.values())
|
|
for k in group_occurrences:
|
|
group_occurrences[k] /= tot
|
|
for k in dataset_occurrences:
|
|
dataset_occurrences[k] /= tot
|
|
|
|
def almost(number):
|
|
return pytest.approx(number, abs=0.02)
|
|
|
|
assert group_occurrences["G1"] == almost(0.2) # group weight: 0.2
|
|
assert group_occurrences["G2"] == almost(0.8) # group weight: 0.8
|
|
assert dataset_occurrences["D1"] == almost(0.1) # group weight: 0.2 * dataset weight 0.5 => 0.1
|
|
assert dataset_occurrences["D2"] == almost(0.1) # group weight: 0.2 * dataset weight 0.5 => 0.1
|
|
assert dataset_occurrences["D3"] == almost(0.4) # group weight: 0.8 * dataset weight 0.5 => 0.4
|
|
assert dataset_occurrences["D4"] == almost(0.4) # group weight: 0.8 * dataset weight 0.5 => 0.4
|
|
|
|
|
|
def test_extended_data_input_cfg_yaml_path(tmp_path, cutset_shar_path, nemo_tarred_manifest_path_multi):
|
|
input_cfg = [
|
|
{
|
|
"type": "nemo_tarred",
|
|
"manifest_filepath": str(nemo_tarred_manifest_path_multi[0]),
|
|
"tarred_audio_filepaths": str(nemo_tarred_manifest_path_multi[1]),
|
|
"weight": 0.5,
|
|
"tags": {
|
|
"language": "en",
|
|
"modality": "audio",
|
|
"dataset_name": "D1",
|
|
},
|
|
},
|
|
{
|
|
"type": "lhotse_shar",
|
|
"shar_path": str(cutset_shar_path),
|
|
"weight": 0.5,
|
|
"tags": {
|
|
"language": "en",
|
|
"modality": "audio",
|
|
"dataset_name": "D2",
|
|
},
|
|
},
|
|
]
|
|
|
|
yaml_path = tmp_path / "input_cfg.yaml"
|
|
lhotse.serialization.save_to_yaml(input_cfg, yaml_path)
|
|
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": input_cfg,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"num_workers": 0,
|
|
"batch_size": 32,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
batch = next(iter(dl))
|
|
assert isinstance(batch, lhotse.CutSet)
|
|
for cut in batch:
|
|
assert cut.dataset_name in ("D1", "D2")
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def txt_en_path(tmp_path_factory):
|
|
tmp_path = tmp_path_factory.mktemp("text_data")
|
|
en_path = tmp_path / "text.en"
|
|
en_path.write_text(
|
|
"""Example text in English.
|
|
Another sentence.
|
|
"""
|
|
)
|
|
return en_path
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def txt_es_path(tmp_path_factory):
|
|
tmp_path = tmp_path_factory.mktemp("text_data")
|
|
es_path = tmp_path / "text.es"
|
|
es_path.write_text(
|
|
"""Otro texto en ingles.
|
|
Otra frase."""
|
|
)
|
|
return es_path
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def questions_path(tmp_path_factory) -> str:
|
|
tmpdir = tmp_path_factory.mktemp("questions")
|
|
qp = tmpdir / "questions.txt"
|
|
qp.write_text("translate the following to spanish")
|
|
return str(qp)
|
|
|
|
|
|
def test_text_file_input(txt_en_path, txt_es_path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
"type": "txt",
|
|
"paths": txt_en_path,
|
|
"language": "en",
|
|
},
|
|
],
|
|
"shuffle": True,
|
|
"num_workers": 0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
# Note: this test does not need to pass a tokenizer because we use static batch sizes
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
# Note: we use islice here because the dataloader will be infinite.
|
|
batches = [batch for batch in islice(dl, 2)]
|
|
|
|
b = batches[0]
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert all(isinstance(c, TextExample) for c in b)
|
|
assert all(c.language == "en" for c in b)
|
|
|
|
b = batches[1]
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert all(isinstance(c, TextExample) for c in b)
|
|
assert all(c.language == "en" for c in b)
|
|
|
|
|
|
def test_text_file_pairs_input(txt_en_path, txt_es_path, questions_path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
"type": "txt_pair",
|
|
"source_paths": txt_en_path,
|
|
"target_paths": txt_es_path,
|
|
"questions_path": questions_path,
|
|
"source_language": "en",
|
|
"target_language": "es",
|
|
"questions_language": "en",
|
|
},
|
|
],
|
|
"shuffle": True,
|
|
"num_workers": 0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
# Note: this test does not need to pass a tokenizer because we use static batch sizes
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
# Note: we use islice here because the dataloader will be infinite.
|
|
batches = [batch for batch in islice(dl, 2)]
|
|
|
|
b = batches[0]
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert all(isinstance(c, SourceTargetTextExample) for c in b)
|
|
assert all(c.source.language == "en" for c in b)
|
|
assert all(c.target.language == "es" for c in b)
|
|
|
|
b = batches[1]
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert all(isinstance(c, SourceTargetTextExample) for c in b)
|
|
assert all(c.source.language == "en" for c in b)
|
|
assert all(c.target.language == "es" for c in b)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def txt_pair_paths_shards(tmp_path_factory, txt_en_path, txt_es_path):
|
|
tmp_path = tmp_path_factory.mktemp("text_data_shards")
|
|
|
|
en_text = txt_en_path.read_text().splitlines()
|
|
(tmp_path / "en_0.txt").write_text("\n".join(en_text[:5]))
|
|
(tmp_path / "en_1.txt").write_text("\n".join(en_text[5:]))
|
|
|
|
es_text = txt_es_path.read_text().splitlines()
|
|
(tmp_path / "es_0.txt").write_text("\n".join(es_text[:5]))
|
|
(tmp_path / "es_1.txt").write_text("\n".join(es_text[5:]))
|
|
|
|
return f"{tmp_path}/en__OP_0..1_CL_.txt", f"{tmp_path}/es__OP_0..1_CL_.txt"
|
|
|
|
|
|
def test_text_file_pairs_shards_input(txt_pair_paths_shards: tuple[str, str], questions_path):
|
|
en_paths, es_paths = txt_pair_paths_shards
|
|
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
"type": "txt_pair",
|
|
"source_paths": en_paths,
|
|
"target_paths": es_paths,
|
|
"questions_path": questions_path,
|
|
"source_language": "en",
|
|
"target_language": "es",
|
|
"questions_language": "en",
|
|
},
|
|
],
|
|
"shuffle": True,
|
|
"num_workers": 0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
# Note: this test does not need to pass a tokenizer because we use static batch sizes
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
# Note: we use islice here because the dataloader will be infinite.
|
|
batches = [batch for batch in islice(dl, 2)]
|
|
|
|
b = batches[0]
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert all(isinstance(c, SourceTargetTextExample) for c in b)
|
|
assert all(c.source.language == "en" for c in b)
|
|
assert all(c.target.language == "es" for c in b)
|
|
|
|
b = batches[1]
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert all(isinstance(c, SourceTargetTextExample) for c in b)
|
|
assert all(c.source.language == "en" for c in b)
|
|
assert all(c.target.language == "es" for c in b)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def en_es_tokenizer(tmp_path_factory, txt_en_path, txt_es_path) -> SentencePieceTokenizer:
|
|
tmpdir = tmp_path_factory.mktemp("en_es_tokenizer")
|
|
text_path = tmpdir / "text.txt"
|
|
text_path.write_text(txt_en_path.read_text() + "\n" + txt_es_path.read_text())
|
|
create_spt_model(text_path, vocab_size=128, sample_size=-1, do_lower_case=False, output_dir=str(tmpdir))
|
|
return SentencePieceTokenizer(str(tmpdir / "tokenizer.model"))
|
|
|
|
|
|
def test_multimodal_text_audio_dataloading(
|
|
txt_pair_paths_shards: tuple[str, str],
|
|
nemo_tarred_manifest_path_multi: tuple[str, str],
|
|
en_es_tokenizer: SentencePieceTokenizer,
|
|
questions_path: str,
|
|
):
|
|
en_paths, es_paths = txt_pair_paths_shards
|
|
manifest_filepath, tarred_audio_filepaths = nemo_tarred_manifest_path_multi
|
|
QF, BT = 50, 1024
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
"type": "txt_pair",
|
|
"source_paths": en_paths,
|
|
"target_paths": es_paths,
|
|
"source_language": "en",
|
|
"target_language": "es",
|
|
"questions_path": questions_path,
|
|
"questions_language": "en",
|
|
"tags": {
|
|
"modality": "text",
|
|
},
|
|
},
|
|
{
|
|
"type": "nemo_tarred",
|
|
"manifest_filepath": manifest_filepath,
|
|
"tarred_audio_filepaths": tarred_audio_filepaths,
|
|
"tags": {
|
|
"modality": "audio",
|
|
},
|
|
},
|
|
],
|
|
"shuffle": True,
|
|
"num_workers": 0,
|
|
"use_multimodal_sampling": True,
|
|
"prompt_format": "plain",
|
|
"batch_tokens": BT,
|
|
# How to set token equivalent duration in actual training?
|
|
# assuming fbank frames: 0.01 is the base due to frame shift;
|
|
# + subsampling x8 gives us 0.08
|
|
# assuming discrete audio tokens, with frame rate 50Hz,
|
|
# we'd get 0.02
|
|
# in this test we'll just use 0.1 for simplicity
|
|
"token_equivalent_duration": 0.1,
|
|
"quadratic_factor": QF,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config,
|
|
global_rank=0,
|
|
world_size=1,
|
|
dataset=Identity(),
|
|
tokenizer=en_es_tokenizer,
|
|
)
|
|
|
|
b = next(iter(dl))
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert len(b)
|
|
assert any(isinstance(ex, Cut) for ex in b)
|
|
assert any(isinstance(ex, SourceTargetTextExample) for ex in b)
|
|
# Batch tokens is not exceeded after applying the quadratic factor correction
|
|
assert sum(ex.num_tokens**2 / QF for ex in b) <= BT
|
|
for ex in b:
|
|
if isinstance(ex, Cut):
|
|
assert ex.modality == "audio"
|
|
assert isinstance(ex.load_audio(), np.ndarray)
|
|
assert isinstance(ex.supervisions[0].text, str)
|
|
if isinstance(ex, SourceTargetTextExample):
|
|
assert ex.modality == "text"
|
|
assert ex.source.language == "en"
|
|
assert ex.target.language == "es"
|
|
assert isinstance(ex.source.text, str)
|
|
assert isinstance(ex.target.text, str)
|
|
assert isinstance(ex.question.text, str)
|
|
assert torch.is_tensor(ex.input_ids)
|
|
assert torch.is_tensor(ex.context_ids)
|
|
assert torch.is_tensor(ex.answer_ids)
|
|
assert torch.is_tensor(ex.mask)
|
|
|
|
|
|
def test_multimodal_text_audio_dataloading_zip_strategy(
|
|
txt_pair_paths_shards: tuple[str, str],
|
|
nemo_tarred_manifest_path_multi: tuple[str, str],
|
|
en_es_tokenizer: SentencePieceTokenizer,
|
|
questions_path: str,
|
|
):
|
|
en_paths, es_paths = txt_pair_paths_shards
|
|
manifest_filepath, tarred_audio_filepaths = nemo_tarred_manifest_path_multi
|
|
QF, BT = 50, 64
|
|
config = OmegaConf.create(
|
|
{
|
|
"multi_config": True,
|
|
"sampler_fusion": "zip", # <---- !!! this option is being tested here !!!
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
"shuffle": True,
|
|
"num_workers": 0,
|
|
"audio": {
|
|
"input_cfg": [
|
|
{
|
|
"type": "nemo_tarred",
|
|
"manifest_filepath": manifest_filepath,
|
|
"tarred_audio_filepaths": tarred_audio_filepaths,
|
|
"tags": {
|
|
"modality": "audio",
|
|
},
|
|
},
|
|
],
|
|
"prompt_format": "plain",
|
|
"use_multimodal_sampling": True,
|
|
"batch_tokens": BT,
|
|
# How to set token equivalent duration in actual training?
|
|
# assuming fbank frames: 0.01 is the base due to frame shift;
|
|
# + subsampling x8 gives us 0.08
|
|
# assuming discrete audio tokens, with frame rate 50Hz,
|
|
# we'd get 0.02
|
|
# in this test we'll just use 0.1 for simplicity
|
|
"token_equivalent_duration": 0.1,
|
|
"quadratic_factor": QF,
|
|
},
|
|
"text": {
|
|
"input_cfg": [
|
|
{
|
|
"type": "txt_pair",
|
|
"source_paths": en_paths,
|
|
"target_paths": es_paths,
|
|
"source_language": "en",
|
|
"target_language": "es",
|
|
"questions_path": questions_path,
|
|
"questions_language": "en",
|
|
"tags": {
|
|
"modality": "text",
|
|
},
|
|
},
|
|
],
|
|
"use_multimodal_sampling": True,
|
|
"prompt_format": "plain",
|
|
"batch_tokens": 64,
|
|
# How to set token equivalent duration in actual training?
|
|
# assuming fbank frames: 0.01 is the base due to frame shift;
|
|
# + subsampling x8 gives us 0.08
|
|
# assuming discrete audio tokens, with frame rate 50Hz,
|
|
# we'd get 0.02
|
|
# in this test we'll just use 0.1 for simplicity
|
|
"token_equivalent_duration": 0.1,
|
|
"quadratic_factor": 50,
|
|
},
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config,
|
|
global_rank=0,
|
|
world_size=1,
|
|
dataset=Identity(),
|
|
tokenizer=en_es_tokenizer,
|
|
)
|
|
|
|
assert isinstance(dl.dataset.sampler, ZipSampler)
|
|
|
|
# Note: we use islice here because the dataloader will be infinite.
|
|
batches = [batch for batch in islice(dl, 2)]
|
|
|
|
b = batches[0]
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert len(b)
|
|
assert any(isinstance(ex, Cut) for ex in b)
|
|
assert any(isinstance(ex, SourceTargetTextExample) for ex in b)
|
|
# Batch tokens is not exceeded after applying the quadratic factor correction
|
|
# Note: zip samples stitches together two batches hence * 2
|
|
assert sum(ex.num_tokens**2 / QF for ex in b) <= BT * 2
|
|
for ex in b:
|
|
if isinstance(ex, Cut):
|
|
assert ex.modality == "audio"
|
|
assert isinstance(ex.load_audio(), np.ndarray)
|
|
assert isinstance(ex.supervisions[0].text, str)
|
|
if isinstance(ex, SourceTargetTextExample):
|
|
assert ex.modality == "text"
|
|
assert ex.source.language == "en"
|
|
assert ex.target.language == "es"
|
|
assert torch.is_tensor(ex.input_ids)
|
|
assert torch.is_tensor(ex.context_ids)
|
|
assert torch.is_tensor(ex.answer_ids)
|
|
assert torch.is_tensor(ex.mask)
|
|
|
|
b = batches[1]
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert len(b)
|
|
assert any(isinstance(ex, Cut) for ex in b)
|
|
assert any(isinstance(ex, SourceTargetTextExample) for ex in b)
|
|
# Batch tokens is not exceeded after applying the quadratic factor correction
|
|
# Note: zip samples stitches together two batches hence * 2
|
|
assert sum(ex.num_tokens**2 / QF for ex in b) <= BT * 2
|
|
for ex in b:
|
|
if isinstance(ex, Cut):
|
|
assert ex.modality == "audio"
|
|
assert isinstance(ex.load_audio(), np.ndarray)
|
|
assert isinstance(ex.supervisions[0].text, str)
|
|
if isinstance(ex, SourceTargetTextExample):
|
|
assert ex.modality == "text"
|
|
assert ex.source.language == "en"
|
|
assert ex.target.language == "es"
|
|
assert torch.is_tensor(ex.input_ids)
|
|
assert torch.is_tensor(ex.context_ids)
|
|
assert torch.is_tensor(ex.answer_ids)
|
|
assert torch.is_tensor(ex.mask)
|
|
|
|
|
|
def test_multimodal_text_audio_dataloading_round_robin_strategy(
|
|
txt_pair_paths_shards: tuple[str, str],
|
|
nemo_tarred_manifest_path_multi: tuple[str, str],
|
|
en_es_tokenizer: SentencePieceTokenizer,
|
|
questions_path: str,
|
|
):
|
|
en_paths, es_paths = txt_pair_paths_shards
|
|
manifest_filepath, tarred_audio_filepaths = nemo_tarred_manifest_path_multi
|
|
QF, BT = 50, 64
|
|
config = OmegaConf.create(
|
|
{
|
|
"multi_config": True,
|
|
"sampler_fusion": "round_robin", # <---- !!! this option is being tested here !!!
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
"shuffle": True,
|
|
"num_workers": 0,
|
|
"audio": {
|
|
"input_cfg": [
|
|
{
|
|
"type": "nemo_tarred",
|
|
"manifest_filepath": manifest_filepath,
|
|
"tarred_audio_filepaths": tarred_audio_filepaths,
|
|
"tags": {
|
|
"modality": "audio",
|
|
},
|
|
},
|
|
],
|
|
"use_multimodal_sampling": True,
|
|
"prompt_format": "plain",
|
|
"batch_tokens": BT,
|
|
# How to set token equivalent duration in actual training?
|
|
# assuming fbank frames: 0.01 is the base due to frame shift;
|
|
# + subsampling x8 gives us 0.08
|
|
# assuming discrete audio tokens, with frame rate 50Hz,
|
|
# we'd get 0.02
|
|
# in this test we'll just use 0.1 for simplicity
|
|
"token_equivalent_duration": 0.1,
|
|
"quadratic_factor": QF,
|
|
},
|
|
"text": {
|
|
"input_cfg": [
|
|
{
|
|
"type": "txt_pair",
|
|
"source_paths": en_paths,
|
|
"target_paths": es_paths,
|
|
"source_language": "en",
|
|
"target_language": "es",
|
|
"questions_path": questions_path,
|
|
"questions_language": "en",
|
|
"tags": {
|
|
"modality": "text",
|
|
},
|
|
},
|
|
],
|
|
"prompt_format": "plain",
|
|
"use_multimodal_sampling": True,
|
|
"batch_tokens": BT,
|
|
# How to set token equivalent duration in actual training?
|
|
# assuming fbank frames: 0.01 is the base due to frame shift;
|
|
# + subsampling x8 gives us 0.08
|
|
# assuming discrete audio tokens, with frame rate 50Hz,
|
|
# we'd get 0.02
|
|
# in this test we'll just use 0.1 for simplicity
|
|
"token_equivalent_duration": 0.1,
|
|
"quadratic_factor": QF,
|
|
},
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config,
|
|
global_rank=0,
|
|
world_size=1,
|
|
dataset=Identity(),
|
|
tokenizer=en_es_tokenizer,
|
|
)
|
|
|
|
assert isinstance(dl.dataset.sampler, RoundRobinSampler)
|
|
|
|
# Note: we use islice here because the dataloader will be infinite.
|
|
batches = [batch for batch in islice(dl, 2)]
|
|
|
|
# Batch 0 is audio-only
|
|
b = batches[0]
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert len(b)
|
|
assert all(isinstance(ex, Cut) for ex in b)
|
|
# Batch tokens is not exceeded after applying the quadratic factor correction
|
|
assert sum(ex.num_tokens**2 / QF for ex in b) <= BT
|
|
for ex in b:
|
|
assert ex.modality == "audio"
|
|
assert isinstance(ex.load_audio(), np.ndarray)
|
|
assert isinstance(ex.supervisions[0].text, str)
|
|
|
|
# Batch 1 is text-only
|
|
b = batches[1]
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert len(b)
|
|
assert all(isinstance(ex, SourceTargetTextExample) for ex in b)
|
|
# Batch tokens is not exceeded after applying the quadratic factor correction
|
|
assert sum(ex.num_tokens**2 / QF for ex in b) <= BT
|
|
for ex in b:
|
|
assert ex.modality == "text"
|
|
assert ex.source.language == "en"
|
|
assert ex.target.language == "es"
|
|
assert torch.is_tensor(ex.input_ids)
|
|
assert torch.is_tensor(ex.context_ids)
|
|
assert torch.is_tensor(ex.answer_ids)
|
|
assert torch.is_tensor(ex.mask)
|
|
|
|
|
|
def test_multimodal_text_audio_dataloading_randomized_round_robin_strategy(
|
|
deterministic_rng,
|
|
txt_pair_paths_shards: tuple[str, str],
|
|
nemo_tarred_manifest_path_multi: tuple[str, str],
|
|
en_es_tokenizer: SentencePieceTokenizer,
|
|
questions_path: str,
|
|
):
|
|
en_paths, es_paths = txt_pair_paths_shards
|
|
manifest_filepath, tarred_audio_filepaths = nemo_tarred_manifest_path_multi
|
|
QF, BT = 50, 64
|
|
config = OmegaConf.create(
|
|
{
|
|
"multi_config": True,
|
|
"sampler_fusion": "randomized_round_robin", # <---- !!! this option is being tested here !!!
|
|
"sampler_weights": {
|
|
"audio": 0.5,
|
|
"text": 0.5,
|
|
},
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
"shuffle": True,
|
|
"num_workers": 0,
|
|
"audio": {
|
|
"input_cfg": [
|
|
{
|
|
"type": "nemo_tarred",
|
|
"manifest_filepath": manifest_filepath,
|
|
"tarred_audio_filepaths": tarred_audio_filepaths,
|
|
"tags": {
|
|
"modality": "audio",
|
|
},
|
|
},
|
|
],
|
|
"use_multimodal_sampling": True,
|
|
"prompt_format": "plain",
|
|
"batch_tokens": BT,
|
|
# How to set token equivalent duration in actual training?
|
|
# assuming fbank frames: 0.01 is the base due to frame shift;
|
|
# + subsampling x8 gives us 0.08
|
|
# assuming discrete audio tokens, with frame rate 50Hz,
|
|
# we'd get 0.02
|
|
# in this test we'll just use 0.1 for simplicity
|
|
"token_equivalent_duration": 0.1,
|
|
"quadratic_factor": QF,
|
|
},
|
|
"text": {
|
|
"input_cfg": [
|
|
{
|
|
"type": "txt_pair",
|
|
"source_paths": en_paths,
|
|
"target_paths": es_paths,
|
|
"source_language": "en",
|
|
"target_language": "es",
|
|
"questions_path": questions_path,
|
|
"questions_language": "en",
|
|
"tags": {
|
|
"modality": "text",
|
|
},
|
|
},
|
|
],
|
|
"prompt_format": "plain",
|
|
"use_multimodal_sampling": True,
|
|
"batch_tokens": BT,
|
|
# How to set token equivalent duration in actual training?
|
|
# assuming fbank frames: 0.01 is the base due to frame shift;
|
|
# + subsampling x8 gives us 0.08
|
|
# assuming discrete audio tokens, with frame rate 50Hz,
|
|
# we'd get 0.02
|
|
# in this test we'll just use 0.1 for simplicity
|
|
"token_equivalent_duration": 0.1,
|
|
"quadratic_factor": QF,
|
|
},
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config,
|
|
global_rank=0,
|
|
world_size=1,
|
|
dataset=Identity(),
|
|
tokenizer=en_es_tokenizer,
|
|
)
|
|
|
|
assert isinstance(dl.dataset.sampler, RoundRobinSampler)
|
|
|
|
# Note: we use islice here because the dataloader will be infinite.
|
|
batches = [batch for batch in islice(dl, 2)]
|
|
|
|
# Batch 0 is audio-only
|
|
b = batches[0]
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert len(b)
|
|
assert all(isinstance(ex, Cut) for ex in b)
|
|
# Batch tokens is not exceeded after applying the quadratic factor correction
|
|
assert sum(ex.num_tokens**2 / QF for ex in b) <= BT
|
|
for ex in b:
|
|
assert ex.modality == "audio"
|
|
assert isinstance(ex.load_audio(), np.ndarray)
|
|
assert isinstance(ex.supervisions[0].text, str)
|
|
|
|
# Batch 1 is text-only
|
|
b = batches[1]
|
|
assert isinstance(b, lhotse.CutSet)
|
|
assert len(b)
|
|
assert all(isinstance(ex, SourceTargetTextExample) for ex in b)
|
|
# Batch tokens is not exceeded after applying the quadratic factor correction
|
|
assert sum(ex.num_tokens**2 / QF for ex in b) <= BT
|
|
for ex in b:
|
|
assert ex.modality == "text"
|
|
assert ex.source.language == "en"
|
|
assert ex.target.language == "es"
|
|
assert torch.is_tensor(ex.input_ids)
|
|
assert torch.is_tensor(ex.context_ids)
|
|
assert torch.is_tensor(ex.answer_ids)
|
|
assert torch.is_tensor(ex.mask)
|
|
|
|
|
|
def test_dataloader_with_noise_nemo_json(cutset_path: Path, nemo_manifest_path: Path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": str(cutset_path),
|
|
"noise_path": str(nemo_manifest_path),
|
|
"noise_mix_prob": 1.0,
|
|
"noise_snr": [-5.0, 5.0],
|
|
"batch_size": 2,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config,
|
|
global_rank=0,
|
|
world_size=1,
|
|
dataset=Identity(),
|
|
)
|
|
batch = next(iter(dl))
|
|
assert isinstance(batch, CutSet)
|
|
assert len(batch) == 2
|
|
cut = batch[0]
|
|
assert isinstance(cut, MixedCut)
|
|
assert -5.0 < cut.tracks[1].snr < 5.0
|
|
cut = batch[1]
|
|
assert isinstance(cut, MixedCut)
|
|
assert -5.0 < cut.tracks[1].snr < 5.0
|
|
|
|
|
|
def test_dataloader_with_noise_lhotse_jsonl(cutset_path: Path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": str(cutset_path),
|
|
"noise_path": str(cutset_path),
|
|
"noise_mix_prob": 1.0,
|
|
"noise_snr": [-5.0, 5.0],
|
|
"batch_size": 2,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config,
|
|
global_rank=0,
|
|
world_size=1,
|
|
dataset=Identity(),
|
|
)
|
|
batch = next(iter(dl))
|
|
assert isinstance(batch, CutSet)
|
|
assert len(batch) == 2
|
|
cut = batch[0]
|
|
assert isinstance(cut, MixedCut)
|
|
assert -5.0 < cut.tracks[1].snr < 5.0
|
|
cut = batch[1]
|
|
assert isinstance(cut, MixedCut)
|
|
assert -5.0 < cut.tracks[1].snr < 5.0
|
|
|
|
|
|
def test_dataloader_with_noise_nemo_tar(cutset_path: Path, nemo_tarred_manifest_path_multi: Path):
|
|
noise_json, noise_tar = nemo_tarred_manifest_path_multi
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": str(cutset_path),
|
|
"noise_path": {
|
|
"manifest_filepath": noise_json,
|
|
"tarred_audio_filepaths": noise_tar,
|
|
},
|
|
"noise_mix_prob": 1.0,
|
|
"noise_snr": [-5.0, 5.0],
|
|
"batch_size": 2,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config,
|
|
global_rank=0,
|
|
world_size=1,
|
|
dataset=Identity(),
|
|
)
|
|
batch = next(iter(dl))
|
|
assert isinstance(batch, CutSet)
|
|
assert len(batch) == 2
|
|
cut = batch[0]
|
|
assert isinstance(cut, MixedCut)
|
|
assert -5.0 < cut.tracks[1].snr < 5.0
|
|
cut = batch[1]
|
|
assert isinstance(cut, MixedCut)
|
|
assert -5.0 < cut.tracks[1].snr < 5.0
|
|
|
|
|
|
def test_dataloader_with_synth_rir(cutset_path: Path):
|
|
from lhotse.augmentation import ReverbWithImpulseResponse
|
|
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": str(cutset_path),
|
|
"rir_enabled": True,
|
|
"rir_prob": 0.5,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config,
|
|
global_rank=0,
|
|
world_size=1,
|
|
dataset=Identity(),
|
|
)
|
|
batch = next(iter(dl))
|
|
assert isinstance(batch, CutSet)
|
|
assert len(batch) == 4
|
|
cut = batch[0]
|
|
assert isinstance(cut, MonoCut)
|
|
assert cut.recording.transforms is None
|
|
cut = batch[1]
|
|
assert isinstance(cut, MonoCut)
|
|
assert cut.recording.transforms is None
|
|
cut = batch[2]
|
|
assert isinstance(cut, MonoCut)
|
|
assert isinstance(cut.recording.transforms, list) and len(cut.recording.transforms) == 1
|
|
tfnm = cut.recording.transforms[0]
|
|
if isinstance(tfnm, dict): # lhotse<=1.23.0
|
|
assert tfnm["name"] == "ReverbWithImpulseResponse"
|
|
else: # lhotse>=1.24.0
|
|
assert isinstance(tfnm, ReverbWithImpulseResponse)
|
|
cut = batch[3]
|
|
assert isinstance(cut, MonoCut)
|
|
assert isinstance(cut.recording.transforms, list) and len(cut.recording.transforms) == 1
|
|
tfnm = cut.recording.transforms[0]
|
|
if isinstance(tfnm, dict): # lhotse<=1.23.0
|
|
assert tfnm["name"] == "ReverbWithImpulseResponse"
|
|
else: # lhotse>=1.24.0
|
|
assert isinstance(tfnm, ReverbWithImpulseResponse)
|
|
|
|
|
|
def test_dataloader_bucket_batch_size(nemo_tarred_manifest_path_multi: tuple[str, str]):
|
|
json_mft, tar_mft = nemo_tarred_manifest_path_multi
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": json_mft,
|
|
"tarred_audio_filepaths": tar_mft,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
# lhotse specific
|
|
"use_bucketing": True,
|
|
"concurrent_bucketing": False,
|
|
# Note: all input cuts belong to the first bucket so the batch size will always be 2.
|
|
"bucket_duration_bins": [2.0, 4.0],
|
|
"bucket_batch_size": [2, 1],
|
|
"drop_last": False,
|
|
"shuffle_buffer_size": 10,
|
|
"bucket_buffer_size": 100,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
for b in islice(dl, 10):
|
|
assert len(b) == 2
|
|
|
|
|
|
@pytest.mark.parametrize("clipping_prob_hard", [0.0, 1.0])
|
|
@pytest.mark.parametrize("clipping_oversampling", [None, 1, 2])
|
|
def test_dataloader_with_clipping_lhotse_jsonl(
|
|
cutset_path: Path, clipping_prob_hard: float, clipping_oversampling: Optional[int]
|
|
):
|
|
from lhotse.augmentation import Clipping, Resample
|
|
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": str(cutset_path),
|
|
"clipping_enabled": True,
|
|
"clipping_gain_db": (0.0, 24.0),
|
|
"clipping_prob": 1.0,
|
|
"clipping_prob_hard": clipping_prob_hard,
|
|
"clipping_oversampling": clipping_oversampling,
|
|
"batch_size": 2,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config,
|
|
global_rank=0,
|
|
world_size=1,
|
|
dataset=Identity(),
|
|
)
|
|
batch = next(iter(dl))
|
|
assert isinstance(batch, CutSet)
|
|
assert len(batch) == 2
|
|
cut = batch[0]
|
|
assert isinstance(cut, MonoCut)
|
|
if clipping_oversampling is not None:
|
|
assert isinstance(cut.recording.transforms[-3], Resample)
|
|
assert isinstance(cut.recording.transforms[-2], Clipping)
|
|
assert isinstance(cut.recording.transforms[-1], Resample)
|
|
else:
|
|
assert isinstance(cut.recording.transforms[-1], Clipping)
|
|
cut = batch[1]
|
|
assert isinstance(cut, MonoCut)
|
|
if clipping_oversampling is not None:
|
|
assert isinstance(cut.recording.transforms[-3], Resample)
|
|
assert isinstance(cut.recording.transforms[-2], Clipping)
|
|
assert isinstance(cut.recording.transforms[-1], Resample)
|
|
else:
|
|
assert isinstance(cut.recording.transforms[-1], Clipping)
|
|
for cut in batch:
|
|
cut.load_audio()
|
|
|
|
|
|
def test_dataloader_with_compression_lossy_lhotse_jsonl(cutset_path: Path):
|
|
from lhotse.augmentation import Compress
|
|
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": str(cutset_path),
|
|
"compression_enabled": True,
|
|
"compression_codecs": ["opus", "mp3", "vorbis"],
|
|
"compression_prob": 1.0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config,
|
|
global_rank=0,
|
|
world_size=1,
|
|
dataset=Identity(),
|
|
)
|
|
batch = next(iter(dl))
|
|
assert isinstance(batch, CutSet)
|
|
assert len(batch) == 4
|
|
cut = batch[0]
|
|
assert isinstance(cut, MonoCut)
|
|
assert isinstance(cut.recording.transforms[-1], Compress)
|
|
cut = batch[1]
|
|
assert isinstance(cut, MonoCut)
|
|
assert isinstance(cut.recording.transforms[-1], Compress)
|
|
for cut in batch:
|
|
cut.load_audio()
|
|
|
|
|
|
def test_dataloader_with_compression_gsm_lhotse_jsonl(cutset_path: Path):
|
|
from lhotse.augmentation import Compress, Resample
|
|
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": str(cutset_path),
|
|
"compression_enabled": True,
|
|
"compression_codecs": ["gsm"],
|
|
"compression_prob": 1.0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config,
|
|
global_rank=0,
|
|
world_size=1,
|
|
dataset=Identity(),
|
|
)
|
|
batch = next(iter(dl))
|
|
assert isinstance(batch, CutSet)
|
|
assert len(batch) == 4
|
|
cut = batch[0]
|
|
assert isinstance(cut, MonoCut)
|
|
assert isinstance(cut.recording.transforms[-3], Resample)
|
|
assert isinstance(cut.recording.transforms[-2], Compress)
|
|
assert isinstance(cut.recording.transforms[-1], Resample)
|
|
for cut in batch:
|
|
cut.load_audio()
|
|
|
|
|
|
def test_dataloader_with_lowpass_using_resampling_lhotse_jsonl(cutset_path: Path):
|
|
from lhotse.augmentation import Resample
|
|
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": str(cutset_path),
|
|
"lowpass_enabled": True,
|
|
"lowpass_frequencies_interval": [3500.0, 4000.0],
|
|
"lowpass_prob": 1.0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config,
|
|
global_rank=0,
|
|
world_size=1,
|
|
dataset=Identity(),
|
|
)
|
|
batch = next(iter(dl))
|
|
assert isinstance(batch, CutSet)
|
|
assert len(batch) == 4
|
|
cut = batch[0]
|
|
assert isinstance(cut, MonoCut)
|
|
assert isinstance(cut.recording.transforms[-2], Resample)
|
|
assert isinstance(cut.recording.transforms[-1], Resample)
|
|
cut = batch[1]
|
|
assert isinstance(cut, MonoCut)
|
|
assert isinstance(cut.recording.transforms[-2], Resample)
|
|
assert isinstance(cut.recording.transforms[-1], Resample)
|
|
for cut in batch:
|
|
cut.load_audio()
|
|
|
|
|
|
def test_dataloader_with_multiple_augmentations_lhotse_jsonl(cutset_path: Path):
|
|
from lhotse.augmentation import Compress, Resample, ReverbWithImpulseResponse
|
|
|
|
config = OmegaConf.create(
|
|
{
|
|
"cuts_path": str(cutset_path),
|
|
"noise_path": str(cutset_path),
|
|
"noise_mix_prob": 1.0,
|
|
"noise_snr": [-5.0, 5.0],
|
|
"rir_enabled": True,
|
|
"rir_prob": 1.0,
|
|
"lowpass_enabled": True,
|
|
"lowpass_frequencies_interval": [3500.0, 4000.0],
|
|
"lowpass_prob": 1.0,
|
|
"compression_enabled": True,
|
|
"compression_codecs": ["gsm"],
|
|
"compression_prob": 1.0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config,
|
|
global_rank=0,
|
|
world_size=1,
|
|
dataset=Identity(),
|
|
)
|
|
batch = next(iter(dl))
|
|
assert isinstance(batch, CutSet)
|
|
assert len(batch) == 4
|
|
cut = batch[0]
|
|
assert isinstance(cut, MixedCut)
|
|
for track in cut.tracks:
|
|
assert isinstance(track.cut.recording.transforms[-3], Resample)
|
|
assert isinstance(track.cut.recording.transforms[-2], Compress)
|
|
assert isinstance(track.cut.recording.transforms[-1], Resample)
|
|
for cut in batch:
|
|
audio = cut.load_audio()
|
|
|
|
|
|
def test_dataloader_2d_bucketing(nemo_tarred_manifest_path_multi: tuple[str, str], en_es_tokenizer):
|
|
json_mft, tar_mft = nemo_tarred_manifest_path_multi
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": json_mft,
|
|
"tarred_audio_filepaths": tar_mft,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
# lhotse specific
|
|
"use_bucketing": True,
|
|
"concurrent_bucketing": False,
|
|
# Here each bin has the format: [audio_duration, token_sequence_length]
|
|
"bucket_duration_bins": [[0.5, 1], [0.5, 2], [2.0, 5], [2.0, 15], [4.0, 10], [4.0, 20]],
|
|
"bucket_batch_size": [7, 6, 5, 4, 3, 2],
|
|
"drop_last": False,
|
|
"shuffle_buffer_size": 10,
|
|
"bucket_buffer_size": 100,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=Identity(), tokenizer=en_es_tokenizer
|
|
)
|
|
|
|
# All of our data have duration 1.0 and 10 tokens so they will fall to bin[3] with batch_size=4
|
|
for b in islice(dl, 10):
|
|
assert len(b) == 4
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def questions_path(tmp_path_factory) -> Path:
|
|
"""A text file with 10 lines containing question values"""
|
|
qdir = tmp_path_factory.mktemp("questions")
|
|
path = qdir / "questions.txt"
|
|
path.write_text("\n".join(f"some question number {i}" for i in range(10)))
|
|
return path
|
|
|
|
|
|
def test_dataloader_from_nemo_nontarred_manifest_with_extra_questions_field_iter(
|
|
nemo_manifest_path: Path, questions_path: Path
|
|
):
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
"manifest_filepath": nemo_manifest_path,
|
|
"type": "nemo",
|
|
"extra_fields": [
|
|
{
|
|
"type": "text_iter",
|
|
"name": "question",
|
|
"path": questions_path,
|
|
}
|
|
],
|
|
},
|
|
],
|
|
"sample_rate": 16000,
|
|
"shuffle": False,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 2,
|
|
"use_bucketing": False,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
b = next(iter(dl))
|
|
c = b[0]
|
|
assert isinstance(c, MonoCut)
|
|
assert hasattr(c, "question")
|
|
assert c.question == "some question number 0"
|
|
c = b[1]
|
|
assert isinstance(c, MonoCut)
|
|
assert hasattr(c, "question")
|
|
assert c.question == "some question number 1"
|
|
|
|
|
|
def test_dataloader_from_nemo_manifest_with_extra_questions_field_iter(
|
|
nemo_tarred_manifest_path: tuple, questions_path: Path
|
|
):
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
"manifest_filepath": nemo_tarred_manifest_path[0],
|
|
"tarred_audio_filepaths": nemo_tarred_manifest_path[1],
|
|
"type": "nemo_tarred",
|
|
"extra_fields": [
|
|
{
|
|
"type": "text_iter",
|
|
"name": "question",
|
|
"path": questions_path,
|
|
}
|
|
],
|
|
},
|
|
],
|
|
"sample_rate": 16000,
|
|
"shuffle": False,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 2,
|
|
"use_bucketing": False,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
b = next(iter(dl))
|
|
c = b[0]
|
|
assert isinstance(c, MonoCut)
|
|
assert hasattr(c, "question")
|
|
assert c.question == "some question number 0"
|
|
c = b[1]
|
|
assert isinstance(c, MonoCut)
|
|
assert hasattr(c, "question")
|
|
assert c.question == "some question number 1"
|
|
|
|
|
|
def test_dataloader_from_nemo_manifest_with_extra_questions_field_sample(
|
|
nemo_tarred_manifest_path: tuple, questions_path: Path
|
|
):
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
"manifest_filepath": nemo_tarred_manifest_path[0],
|
|
"tarred_audio_filepaths": nemo_tarred_manifest_path[1],
|
|
"type": "nemo_tarred",
|
|
"extra_fields": [
|
|
{
|
|
"type": "text_sample",
|
|
"name": "question",
|
|
"path": questions_path,
|
|
}
|
|
],
|
|
},
|
|
],
|
|
"sample_rate": 16000,
|
|
"shuffle": False,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 5,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
"use_bucketing": False,
|
|
}
|
|
)
|
|
|
|
# Note: despite shuffle=True, it is sampling lines from questions_path because of type: "text_sample"
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
b = next(iter(dl))
|
|
c = b[0]
|
|
assert isinstance(c, MonoCut)
|
|
assert hasattr(c, "question")
|
|
assert c.question == "some question number 6"
|
|
c = b[1]
|
|
assert isinstance(c, MonoCut)
|
|
assert hasattr(c, "question")
|
|
assert c.question == "some question number 6"
|
|
c = b[2]
|
|
assert isinstance(c, MonoCut)
|
|
assert hasattr(c, "question")
|
|
assert c.question == "some question number 0"
|
|
c = b[3]
|
|
assert isinstance(c, MonoCut)
|
|
assert hasattr(c, "question")
|
|
assert c.question == "some question number 4"
|
|
c = b[4]
|
|
assert isinstance(c, MonoCut)
|
|
assert hasattr(c, "question")
|
|
assert c.question == "some question number 8"
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def nemo_tarred_manifest_path_with_offset(tmp_path_factory) -> Tuple[str, str]:
|
|
"""10 utterances of length 1s as a NeMo tarred manifest."""
|
|
from lhotse.serialization import SequentialJsonlWriter
|
|
from lhotse.shar.writers import TarWriter
|
|
|
|
root = tmp_path_factory.mktemp("nemo_tar_offset")
|
|
root.mkdir(exist_ok=True)
|
|
recording = dummy_recording(0, duration=10.0, with_data=True)
|
|
|
|
with (
|
|
TarWriter(f"{root}/audios_0.tar", shard_size=None) as tar_writer,
|
|
SequentialJsonlWriter(root / "tarred_audio_filepaths.jsonl") as mft_writer,
|
|
):
|
|
|
|
def audio_path(n: int = None):
|
|
return recording.id + ("" if n is None else f"-sub{n}") + ".wav"
|
|
|
|
tar_writer.write(audio_path(), BytesIO(recording.sources[0].source))
|
|
mft_writer.write(
|
|
{ # segment 0-3s
|
|
"audio_filepath": audio_path(),
|
|
"offset": 0.0,
|
|
"duration": 3.0,
|
|
"text": "irrelevant",
|
|
"lang": "en",
|
|
"shard_id": 0,
|
|
}
|
|
)
|
|
mft_writer.write(
|
|
{ # segment 4-9s
|
|
"audio_filepath": audio_path(1),
|
|
"offset": 4.0,
|
|
"duration": 5.0,
|
|
"text": "irrelevant-2",
|
|
"lang": "en",
|
|
"shard_id": 0,
|
|
}
|
|
)
|
|
mft_writer.write(
|
|
{ # full recording - for reference
|
|
"audio_filepath": audio_path(2),
|
|
"offset": 0.0,
|
|
"duration": 10.0,
|
|
"text": "irrelevant irrelevant-2",
|
|
"lang": "en",
|
|
"shard_id": 0,
|
|
}
|
|
)
|
|
return mft_writer.path, tar_writer.output_paths[0]
|
|
|
|
|
|
def test_dataloader_from_tarred_nemo_manifest_with_offset(nemo_tarred_manifest_path_with_offset: tuple[str, str]):
|
|
json_mft, tar_mft = nemo_tarred_manifest_path_with_offset
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": json_mft,
|
|
"tarred_audio_filepaths": tar_mft,
|
|
"sample_rate": 16000,
|
|
"shuffle": False,
|
|
"num_workers": 0,
|
|
"batch_size": 3,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
"force_finite": True,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
# Loads all three examples in a single mini-batch (that's why batch_size=3).
|
|
batches = [b for b in dl]
|
|
assert len(batches) == 1
|
|
(batch,) = batches
|
|
assert len(batch) == 3
|
|
|
|
# Validate example containing full 10s recording.
|
|
full_cut = batch[1]
|
|
assert full_cut.start == 0.0
|
|
assert full_cut.duration == 10.0
|
|
assert full_cut.supervisions[0].text == "irrelevant irrelevant-2"
|
|
assert full_cut.supervisions[0].language == "en"
|
|
full_audio = full_cut.load_audio()
|
|
assert full_audio.shape[1] == full_cut.num_samples == 160000 # 10s * 16kHz
|
|
|
|
# Validate segment 0-3s.
|
|
cut = batch[2]
|
|
assert cut.start == 0.0
|
|
assert cut.duration == 3.0
|
|
assert cut.supervisions[0].text == "irrelevant"
|
|
assert cut.supervisions[0].language == "en"
|
|
audio = cut.load_audio()
|
|
assert audio.shape[1] == cut.num_samples
|
|
# Check the audio for the segment is identical to a slice of the full audio.
|
|
np.testing.assert_equal(audio, full_audio[:, : compute_num_samples(cut.duration, cut.sampling_rate)])
|
|
|
|
# Validate segment 4-9s.
|
|
# Note: LazyNeMoTarredIterator removes the offset information, as it creates a new recording
|
|
# that's a "subset" of the original recording as a memory saving optimization.
|
|
# Hence, we will not see cut.start == 4.0.
|
|
cut = batch[0]
|
|
assert cut.start == 0.0
|
|
assert cut.duration == 5.0
|
|
assert cut.supervisions[0].text == "irrelevant-2"
|
|
assert cut.supervisions[0].language == "en"
|
|
audio = cut.load_audio()
|
|
assert audio.shape[1] == cut.num_samples
|
|
# Check the audio for the segment is identical to a slice of the full audio.
|
|
np.testing.assert_equal(
|
|
audio, full_audio[:, compute_num_samples(4.0, cut.sampling_rate) : compute_num_samples(9.0, cut.sampling_rate)]
|
|
)
|
|
|
|
|
|
def test_force_iterable_dataset(cutset_path: Path):
|
|
config = OmegaConf.create({"cuts_path": cutset_path, "batch_size": 2, "num_workers": 2})
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
batches_map = [b for b in dl]
|
|
|
|
config = OmegaConf.create(
|
|
{"cuts_path": cutset_path, "batch_size": 2, "num_workers": 2, "force_iterable_dataset": True}
|
|
)
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
batches_iter = [b for b in dl]
|
|
|
|
# 2x duplicated data due to iterable dataset lack of deduplication
|
|
assert len(batches_iter) == 2 * len(batches_map)
|
|
# assertion that this is in fact the same data (same ids)
|
|
assert set(c.id for b in batches_iter for c in b) == set(c.id for b in batches_map for c in b)
|
|
|
|
|
|
def test_force_map_dataset(cutset_shar_path: Path):
|
|
config = OmegaConf.create({"shar_path": cutset_shar_path, "batch_size": 2, "num_workers": 2, "force_finite": True})
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
batches_iter = [b for b in dl]
|
|
|
|
config = OmegaConf.create(
|
|
{
|
|
"shar_path": cutset_shar_path,
|
|
"batch_size": 2,
|
|
"num_workers": 2,
|
|
"force_map_dataset": True,
|
|
"force_finite": True,
|
|
}
|
|
)
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
batches_map = [b for b in dl]
|
|
|
|
# 2x duplicated data due to iterable dataset lack of deduplication
|
|
assert len(batches_iter) == 2 * len(batches_map)
|
|
# assertion that this is in fact the same data (same ids)
|
|
assert set(c.id for b in batches_iter for c in b) == set(c.id for b in batches_map for c in b)
|
|
|
|
|
|
def test_dataloader_from_tarred_nemo_subset_manifest(nemo_tarred_manifest_subset_path: tuple[str, str]):
|
|
json_mft, tar_mft, subset_items = nemo_tarred_manifest_subset_path
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": json_mft,
|
|
"tarred_audio_filepaths": tar_mft,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
# lhotse specific
|
|
"use_bucketing": True,
|
|
"concurrent_bucketing": False,
|
|
"num_buckets": 2,
|
|
"drop_last": False,
|
|
"batch_duration": 4.0, # seconds
|
|
"quadratic_duration": 15.0, # seconds
|
|
"shuffle_buffer_size": 10,
|
|
"bucket_buffer_size": 100,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
"tarred_random_access": True,
|
|
"force_finite": True,
|
|
}
|
|
)
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
seen_ids = list()
|
|
for batch in dl:
|
|
current_ids = batch["ids"]
|
|
seen_ids += current_ids
|
|
|
|
expected_ids = set([data['audio_filepath'] for data in subset_items])
|
|
seen_ids_set = set(seen_ids)
|
|
assert len(seen_ids_set) == len(seen_ids), "Duplicate IDs found in the batch."
|
|
assert seen_ids_set == expected_ids, "The set of IDs in the batches does not match the input JSON manifests."
|
|
|
|
|
|
def test_dataloader_from_nemo_manifest_with_skipme(nemo_manifest_with_skipme_path: Path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": nemo_manifest_with_skipme_path,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 1,
|
|
# lhotse specific
|
|
"use_bucketing": False,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=_Identity())
|
|
batches = [batch for batch in dl]
|
|
skipme_s = [cut.custom.get('_skipme', 0) for batch in batches for cut in batch]
|
|
|
|
assert len(batches) == 8
|
|
assert not any(skipme_s)
|
|
|
|
|
|
def test_dataloader_from_tarred_nemo_manifest_with_skipme(nemo_tarred_manifest_with_skipme_path: tuple[Path, str]):
|
|
json_mft, tar_mft = nemo_tarred_manifest_with_skipme_path
|
|
config = OmegaConf.create(
|
|
{
|
|
"manifest_filepath": json_mft,
|
|
"tarred_audio_filepaths": tar_mft,
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 1,
|
|
# lhotse specific
|
|
"use_bucketing": False,
|
|
"force_finite": True,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=_Identity())
|
|
batches = [batch for batch in dl]
|
|
skipme_s = [cut.custom.get('_skipme', 0) for batch in batches for cut in batch]
|
|
|
|
assert len(batches) == 8
|
|
assert not any(skipme_s)
|
|
|
|
|
|
def test_dataloader_from_data_input_cfg_yaml_path_with_skipme(cutset_shar_path, nemo_tarred_manifest_with_skipme_path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
"type": "nemo_tarred",
|
|
"manifest_filepath": nemo_tarred_manifest_with_skipme_path[0],
|
|
"tarred_audio_filepaths": nemo_tarred_manifest_with_skipme_path[1],
|
|
"weight": 0.5,
|
|
"tags": {
|
|
"language": "en",
|
|
"modality": "audio",
|
|
"dataset_name": "D1",
|
|
},
|
|
},
|
|
{
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path,
|
|
"weight": 0.5,
|
|
"tags": {
|
|
"language": "en",
|
|
"modality": "audio",
|
|
"dataset_name": "D2",
|
|
},
|
|
},
|
|
],
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"num_workers": 0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
"force_finite": True,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
batches = [batch for batch in dl]
|
|
skipme_s = [cut.custom.get('_skipme', 0) for batch in batches for cut in batch]
|
|
|
|
assert not any(skipme_s)
|
|
|
|
|
|
def test_dataloader_lhotse_shar_nemo_tarred_slice_length(
|
|
nemo_tarred_manifest_path_multi: tuple[str, str], cutset_shar_path: Path
|
|
):
|
|
json_mft, tar_mft = nemo_tarred_manifest_path_multi
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
# 2 shards, 5 utterances each
|
|
"type": "nemo_tarred",
|
|
"manifest_filepath": json_mft,
|
|
"tarred_audio_filepaths": tar_mft,
|
|
"tags": {"origin": "nemo_tarred"},
|
|
},
|
|
{
|
|
# 2 shards, 5 utterances each
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path,
|
|
"tags": {"origin": "lhotse_shar"},
|
|
},
|
|
],
|
|
"slice_length": 2,
|
|
"shuffle": True,
|
|
"num_workers": 0,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
"batch_size": 4,
|
|
"force_finite": True,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
batches = [b for b in dl]
|
|
assert len(batches) == 2
|
|
|
|
# We expect to sample a total of 8 examples (4 shards, 2 examples each), in 2 batches of size 4 each,
|
|
# half of it coming from nemo tarred dataset, and the other half from lhotse shar dataset.
|
|
origin_tags = Counter()
|
|
origin_shards = Counter()
|
|
for b in batches:
|
|
assert len(b) == 4
|
|
for cut in b:
|
|
origin_tags[cut.origin] += 1
|
|
if cut.origin == "lhotse_shar":
|
|
origin_shard = cut.shard_origin
|
|
else:
|
|
origin_shard = cut.manifest_origin
|
|
origin_shard = Path(origin_shard).name
|
|
origin_shards[origin_shard] += 1
|
|
assert origin_tags["lhotse_shar"] == 4
|
|
assert origin_tags["nemo_tarred"] == 4
|
|
assert origin_shards["cuts.000000.jsonl.gz"] == 2
|
|
assert origin_shards["cuts.000001.jsonl.gz"] == 2
|
|
assert origin_shards["manifest_0.jsonl"] == 2
|
|
assert origin_shards["manifest_1.jsonl"] == 2
|
|
|
|
|
|
def test_dataloader_lhotse_shar_slice_length_multi_epoch_different_sample(cutset_shar_path: Path):
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
# 2 shards, 5 utterances each
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path,
|
|
},
|
|
],
|
|
"slice_length": 2,
|
|
"shuffle": False, # shuffle is disabled, but the slice offset still must be random!
|
|
"num_workers": 0,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
"batch_size": 2,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
# 2 batches == 1 epoch => 4 batches == 2 epochs
|
|
batches = [b for b in islice(dl, 4)]
|
|
assert len(batches) == 4
|
|
epoch0_ids = [cut.id for b in batches[:2] for cut in b]
|
|
epoch1_ids = [cut.id for b in batches[2:] for cut in b]
|
|
assert epoch0_ids != epoch1_ids
|
|
assert epoch0_ids + epoch1_ids != sorted(epoch0_ids + epoch1_ids) # true when slice_length=None
|
|
|
|
|
|
def test_dataloader_nemo_tarred_slice_length_multi_epoch_different_sample(
|
|
nemo_tarred_manifest_path_multi: tuple[str, str],
|
|
):
|
|
json_mft, tar_mft = nemo_tarred_manifest_path_multi
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
# 2 shards, 5 utterances each
|
|
"type": "nemo_tarred",
|
|
"manifest_filepath": json_mft,
|
|
"tarred_audio_filepaths": tar_mft,
|
|
},
|
|
],
|
|
"slice_length": 2,
|
|
"shuffle": False, # shuffle is disabled, but the slice offset still must be random!
|
|
"num_workers": 0,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
"batch_size": 2,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
# 2 batches == 1 epoch => 4 batches == 2 epochs
|
|
batches = [b for b in islice(dl, 4)]
|
|
assert len(batches) == 4
|
|
|
|
epoch0_ids = [cut.id for b in batches[:2] for cut in b]
|
|
epoch1_ids = [cut.id for b in batches[2:] for cut in b]
|
|
assert epoch0_ids != epoch1_ids
|
|
assert epoch0_ids + epoch1_ids != sorted(epoch0_ids + epoch1_ids) # true when slice_length=None
|
|
|
|
|
|
def test_dataloader_reweight_temperature_intermediate_value(
|
|
deterministic_rng, cutset_shar_path: Path, cutset_shar_path_other: Path
|
|
):
|
|
"""
|
|
Test that reweight_temperature=0.5 gives intermediate behavior between equal and original weights.
|
|
With temperature=0.5, datasets with weights 900 and 100:
|
|
- Original (temp=1.0): 90-10
|
|
- Equal (temp=0.0): 50-50
|
|
- Expected (temp=0.5): ~75-25 (since 900^0.5 / (900^0.5 + 100^0.5) ≈ 0.75)
|
|
"""
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{"type": "lhotse_shar", "shar_path": cutset_shar_path, "weight": 900},
|
|
{"type": "lhotse_shar", "shar_path": cutset_shar_path_other, "weight": 100},
|
|
],
|
|
"reweight_temperature": [0.5],
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
# Sample multiple batches and count occurrences
|
|
dataset_counts = Counter()
|
|
for batch in islice(dl, 50):
|
|
for cid in batch["ids"]:
|
|
if cid.startswith("dummy"):
|
|
dataset_counts["dataset1"] += 1
|
|
elif cid.startswith("other"):
|
|
dataset_counts["dataset2"] += 1
|
|
|
|
total = sum(dataset_counts.values())
|
|
# With temperature=0.5, expect approximately 75-25 distribution
|
|
assert dataset_counts["dataset1"] / total == pytest.approx(0.75, abs=0.1)
|
|
assert dataset_counts["dataset2"] / total == pytest.approx(0.25, abs=0.1)
|
|
|
|
|
|
def test_dataloader_reweight_temperature_nested_groups(
|
|
deterministic_rng,
|
|
cutset_shar_path: Path,
|
|
cutset_shar_path_other: Path,
|
|
nemo_tarred_manifest_path_multi: tuple[str, str],
|
|
):
|
|
"""
|
|
Test that reweight_temperature works correctly with nested groups.
|
|
Using [1.0, 0.0]: level 1 preserves weights, level 2 equalizes.
|
|
|
|
Structure:
|
|
- Group A (weight=200):
|
|
- Dataset A1 (weight=180)
|
|
- Dataset A2 (weight=20)
|
|
- Group B (weight=800):
|
|
- Dataset B1 (weight=800)
|
|
|
|
Expected:
|
|
- Level 1 (temp=1.0): Group A gets 20%, Group B gets 80%
|
|
- Level 2 (temp=0.0): Within Group A, A1 and A2 each get 50% (so 10% and 10% of total)
|
|
"""
|
|
json_mft, tar_mft = nemo_tarred_manifest_path_multi
|
|
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
"type": "group",
|
|
"weight": 200,
|
|
"tags": {"group": "A"},
|
|
"input_cfg": [
|
|
{
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path,
|
|
"weight": 180,
|
|
"tags": {"dataset_name": "A1"},
|
|
},
|
|
{
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path_other,
|
|
"weight": 20,
|
|
"tags": {"dataset_name": "A2"},
|
|
},
|
|
],
|
|
},
|
|
{
|
|
"type": "group",
|
|
"weight": 800,
|
|
"tags": {"group": "B"},
|
|
"input_cfg": [
|
|
{
|
|
"type": "nemo_tarred",
|
|
"manifest_filepath": json_mft,
|
|
"tarred_audio_filepaths": tar_mft,
|
|
"weight": 800,
|
|
"tags": {"dataset_name": "B1"},
|
|
},
|
|
],
|
|
},
|
|
],
|
|
"reweight_temperature": [1.0, 0.0], # Level 1: preserve, Level 2: equalize
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
# Sample multiple batches and count occurrences using tags dataset_name
|
|
group_counts = Counter()
|
|
dataset_counts = Counter()
|
|
for batch in islice(dl, 100):
|
|
for cut in batch:
|
|
group_counts[cut.group] += 1
|
|
dataset_counts[cut.dataset_name] += 1
|
|
|
|
total = sum(group_counts.values())
|
|
# Level 1: temperature=1.0, so groups A and B should have 20-80 split
|
|
assert group_counts["A"] / total == pytest.approx(0.2, abs=0.1)
|
|
assert group_counts["B"] / total == pytest.approx(0.8, abs=0.1)
|
|
|
|
# Level 2 (within group A): temperature=0.0, so A1 and A2 should be equal (50-50)
|
|
# which means each should be ~10% of total (0.5 * 0.2)
|
|
if group_counts["A"] > 0: # Make sure we have samples from group A
|
|
a_total = dataset_counts["A1"] + dataset_counts["A2"]
|
|
assert dataset_counts["A1"] / a_total == pytest.approx(0.5, abs=0.15)
|
|
assert dataset_counts["A2"] / a_total == pytest.approx(0.5, abs=0.15)
|
|
|
|
|
|
def test_dataloader_reweight_temperature_three_datasets(
|
|
deterministic_rng,
|
|
cutset_shar_path: Path,
|
|
cutset_shar_path_other: Path,
|
|
nemo_tarred_manifest_path_multi: tuple[str, str],
|
|
):
|
|
"""
|
|
Test reweight_temperature with three datasets of different sizes.
|
|
With temperature=0.0, all three should be sampled equally.
|
|
"""
|
|
json_mft, tar_mft = nemo_tarred_manifest_path_multi
|
|
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{"type": "lhotse_shar", "shar_path": cutset_shar_path, "weight": 600, "tags": {"dataset_name": "D1"}},
|
|
{
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path_other,
|
|
"weight": 300,
|
|
"tags": {"dataset_name": "D2"},
|
|
},
|
|
{
|
|
"type": "nemo_tarred",
|
|
"manifest_filepath": json_mft,
|
|
"tarred_audio_filepaths": tar_mft,
|
|
"weight": 100,
|
|
"tags": {"dataset_name": "D3"},
|
|
},
|
|
],
|
|
"reweight_temperature": [0.0], # Equalize all three
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 6,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
# Sample multiple batches and count occurrences using the dataset_name tag
|
|
dataset_counts = Counter()
|
|
for batch in islice(dl, 50):
|
|
for cut in batch:
|
|
dataset_counts[cut.dataset_name] += 1
|
|
|
|
total = sum(dataset_counts.values())
|
|
# With temperature=0.0, expect approximately equal distribution (33-33-33)
|
|
assert dataset_counts["D1"] / total == pytest.approx(1 / 3, abs=0.1)
|
|
assert dataset_counts["D2"] / total == pytest.approx(1 / 3, abs=0.1)
|
|
assert dataset_counts["D3"] / total == pytest.approx(1 / 3, abs=0.1)
|
|
|
|
|
|
def test_dataloader_reweight_temperature_deeply_nested(
|
|
deterministic_rng,
|
|
cutset_shar_path: Path,
|
|
cutset_shar_path_other: Path,
|
|
nemo_tarred_manifest_path_multi: tuple[str, str],
|
|
):
|
|
"""
|
|
Test reweight_temperature with three levels of nesting using [1.0, 0.5, 0.0].
|
|
This verifies that temperature is correctly applied at each level of the hierarchy.
|
|
"""
|
|
json_mft, tar_mft = nemo_tarred_manifest_path_multi
|
|
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
{
|
|
"type": "group",
|
|
"weight": 700,
|
|
"tags": {"level1": "GroupX"},
|
|
"input_cfg": [
|
|
{
|
|
"type": "group",
|
|
"weight": 400,
|
|
"tags": {"level2": "GroupX1"},
|
|
"input_cfg": [
|
|
{
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path,
|
|
"weight": 300,
|
|
"tags": {"dataset_name": "X1a"},
|
|
},
|
|
{
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path_other,
|
|
"weight": 100,
|
|
"tags": {"dataset_name": "X1b"},
|
|
},
|
|
],
|
|
},
|
|
{
|
|
"type": "group",
|
|
"weight": 300,
|
|
"tags": {"level2": "GroupX2"},
|
|
"input_cfg": [
|
|
{
|
|
"type": "nemo_tarred",
|
|
"manifest_filepath": json_mft,
|
|
"tarred_audio_filepaths": tar_mft,
|
|
"weight": 300,
|
|
"tags": {"dataset_name": "X2a"},
|
|
},
|
|
],
|
|
},
|
|
],
|
|
},
|
|
{
|
|
"type": "group",
|
|
"weight": 300,
|
|
"tags": {"level1": "GroupY"},
|
|
"input_cfg": [
|
|
{
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path,
|
|
"weight": 300,
|
|
"tags": {"dataset_name": "Y1"},
|
|
},
|
|
],
|
|
},
|
|
],
|
|
"reweight_temperature": [1.0, 0.5, 0.0], # Level 1: preserve, Level 2: intermediate, Level 3: equalize
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(
|
|
config=config, global_rank=0, world_size=1, dataset=UnsupervisedAudioDataset()
|
|
)
|
|
|
|
# Sample multiple batches and just verify it works without errors
|
|
# The exact distribution is complex to calculate, so we just ensure the dataloader runs
|
|
batches = [batch for batch in islice(dl, 20)]
|
|
assert len(batches) == 20
|
|
for batch in batches:
|
|
assert "audio" in batch
|
|
assert "audio_lens" in batch
|
|
assert "ids" in batch
|
|
|
|
|
|
def test_dataloader_reweight_temperature_mixed_leaf_and_group(
|
|
deterministic_rng, cutset_shar_path: Path, cutset_shar_path_other: Path
|
|
):
|
|
"""
|
|
Test config with mix of leaf items and groups at the same level.
|
|
The leaf item doesn't have nested input_cfg, but the group does.
|
|
Max depth should be 2, requiring 2 temperatures.
|
|
"""
|
|
config = OmegaConf.create(
|
|
{
|
|
"input_cfg": [
|
|
# Leaf item at level 1 (no nested input_cfg)
|
|
{
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path,
|
|
"weight": 300,
|
|
"tags": {"dataset_name": "Leaf"},
|
|
},
|
|
# Group with nested items
|
|
{
|
|
"type": "group",
|
|
"weight": 700,
|
|
"tags": {"group": "Nested"},
|
|
"input_cfg": [
|
|
{
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path_other,
|
|
"weight": 600,
|
|
"tags": {"dataset_name": "N1"},
|
|
},
|
|
{
|
|
"type": "lhotse_shar",
|
|
"shar_path": cutset_shar_path,
|
|
"weight": 100,
|
|
"tags": {"dataset_name": "N2"},
|
|
},
|
|
],
|
|
},
|
|
],
|
|
"reweight_temperature": [1.0, 0.0], # 2 levels: level 1 preserves, level 2 equalizes
|
|
"sample_rate": 16000,
|
|
"shuffle": True,
|
|
"use_lhotse": True,
|
|
"num_workers": 0,
|
|
"batch_size": 4,
|
|
"seed": 0,
|
|
"shard_seed": 0,
|
|
}
|
|
)
|
|
|
|
dl = get_lhotse_dataloader_from_config(config=config, global_rank=0, world_size=1, dataset=Identity())
|
|
|
|
# Sample multiple batches and count occurrences
|
|
dataset_counts = Counter()
|
|
for batch in islice(dl, 100):
|
|
for cut in batch:
|
|
dataset_counts[cut.dataset_name] += 1
|
|
|
|
total = sum(dataset_counts.values())
|
|
|
|
# Level 1: Leaf gets 30%, Nested group gets 70% (temperature=1.0 preserves weights)
|
|
assert dataset_counts["Leaf"] / total == pytest.approx(0.3, abs=0.1)
|
|
|
|
# Level 2 within Nested group: N1 and N2 should be equalized (temperature=0.0)
|
|
# Each should get ~35% of total (0.7 * 0.5)
|
|
nested_total = dataset_counts["N1"] + dataset_counts["N2"]
|
|
assert dataset_counts["N1"] / nested_total == pytest.approx(0.5, abs=0.15)
|
|
assert dataset_counts["N2"] / nested_total == pytest.approx(0.5, abs=0.15)
|