Spaces:
Runtime error
Runtime error
File size: 3,582 Bytes
0558aa4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
# 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.
import pytest
import torch
from lhotse import CutSet
from lhotse.testing.dummies import DummyManifest
from lightning.pytorch.utilities import CombinedLoader
from omegaconf import DictConfig
from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer, create_spt_model
from nemo.collections.speechlm2.data import DataModule
@pytest.fixture
def data_config(tmp_path):
ap, cp = tmp_path / "audio", str(tmp_path) + "/{tag}_cuts.jsonl.gz"
def _assign(k, v):
def _inner(obj):
setattr(obj, k, v)
return obj
return _inner
for tag in ("train", "val_set_0", "val_set_1"):
(
DummyManifest(CutSet, begin_id=0, end_id=2, with_data=True)
.map(_assign("tag", tag))
.save_audios(ap)
.drop_in_memory_data()
.to_file(cp.format(tag=tag))
)
return DictConfig(
{
"train_ds": {
"input_cfg": [
{
"type": "lhotse",
"cuts_path": cp.format(tag="train"),
}
],
"batch_size": 2,
},
"validation_ds": {
"datasets": {
"val_set_0": {"cuts_path": cp.format(tag="val_set_0")},
"val_set_1": {"cuts_path": cp.format(tag="val_set_1")},
},
"batch_size": 2,
},
}
)
@pytest.fixture
def tokenizer(tmp_path_factory):
tmpdir = tmp_path_factory.mktemp("tok")
text_path = tmpdir / "text.txt"
text_path.write_text("\n".join(chr(i) for i in range(256)))
create_spt_model(
text_path,
vocab_size=512,
sample_size=-1,
do_lower_case=False,
output_dir=str(tmpdir),
bos=True,
eos=True,
remove_extra_whitespaces=True,
)
return SentencePieceTokenizer(str(tmpdir / "tokenizer.model"))
class Identity(torch.utils.data.Dataset):
def __getitem__(self, item):
return item
def test_datamodule_train_dataloader(data_config, tokenizer):
data = DataModule(data_config, tokenizer=tokenizer, dataset=Identity())
dl = data.train_dataloader()
assert isinstance(dl, torch.utils.data.DataLoader)
dli = iter(dl)
batch = next(dli)
assert isinstance(batch, CutSet)
assert len(batch) == 2
assert all(c.tag == "train" for c in batch)
def test_datamodule_validation_dataloader(data_config, tokenizer):
val_sets = {"val_set_0", "val_set_1"}
data = DataModule(data_config, tokenizer=tokenizer, dataset=Identity())
dl = data.val_dataloader()
assert isinstance(dl, CombinedLoader)
dli = iter(dl)
batch, batch_idx, dataloader_idx = next(dli)
assert isinstance(batch, dict)
assert batch.keys() == val_sets
for vs in val_sets:
assert len(batch[vs]) == 2
assert all(c.tag == vs for c in batch[vs])
|