wehub-resource-sync
ba4be087d5
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
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
CICD NeMo / cicd-test-container-build (push) Has been cancelled
CICD NeMo / cicd-import-tests (push) Has been cancelled
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
CICD NeMo / cicd-main-speech (push) Has been cancelled
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
CICD NeMo / Coverage (e2e) (push) Has been cancelled
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
1069 lines
49 KiB
Python
1069 lines
49 KiB
Python
# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. 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 json
|
|
import os
|
|
import random
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import librosa
|
|
import numpy as np
|
|
import torch.utils.data
|
|
|
|
from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
|
|
from nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers import BaseTokenizer, IPABPETokenizer
|
|
from nemo.collections.tts.parts.preprocessing.feature_processors import FeatureProcessor
|
|
from nemo.collections.tts.parts.preprocessing.features import Featurizer
|
|
from nemo.collections.tts.parts.utils.tts_dataset_utils import (
|
|
_read_audio,
|
|
beta_binomial_prior_distribution,
|
|
chunk_text_for_inference,
|
|
filter_dataset_by_duration,
|
|
get_tokenizer_for_language,
|
|
get_weighted_sampler,
|
|
load_audio,
|
|
setup_pronunciation_control_g2p,
|
|
stack_tensors,
|
|
tokenize_text_with_pronunciation_control,
|
|
)
|
|
from nemo.core.classes import Dataset
|
|
from nemo.utils import logging
|
|
|
|
|
|
@dataclass
|
|
class DatasetMeta:
|
|
manifest_path: Path
|
|
audio_dir: Path
|
|
feature_dir: Optional[Path] = None
|
|
sample_weight: float = 1.0
|
|
language: Optional[str] = None
|
|
tokenizer_names: List[str] = None
|
|
|
|
|
|
@dataclass
|
|
class DatasetSample:
|
|
dataset_name: str
|
|
manifest_entry: Dict[str, Any]
|
|
audio_dir: Optional[Path]
|
|
feature_dir: Optional[Path]
|
|
text: str
|
|
language: Optional[str] = None
|
|
speaker: Optional[str] = None
|
|
speaker_index: int = None
|
|
tokenizer_names: List[str] = None
|
|
|
|
|
|
class TextToSpeechDataset(Dataset):
|
|
"""
|
|
Class for processing and loading text to speech training examples.
|
|
|
|
Args:
|
|
dataset_meta: Dict of dataset names (string) to dataset metadata.
|
|
sample_rate: Sample rate to load audio as. If the audio is stored at a different sample rate, then it will
|
|
be resampled.
|
|
text_tokenizer: Tokenizer to apply to the text field.
|
|
weighted_sampling_steps_per_epoch: Optional int, If provided, then data will be sampled (with replacement) based on
|
|
the sample weights provided in the dataset metadata. If None, then sample weights will be ignored.
|
|
speaker_path: Optional, path to JSON file with speaker indices, for multi-speaker training. Can be created with
|
|
scripts.dataset_processing.tts.create_speaker_map.py
|
|
featurizers: Optional, list of featurizers to load feature data from. Should be the same config provided
|
|
when running scripts.dataset_processing.tts.compute_features.py before training.
|
|
feature_processors: Optional, list of feature processors to run on training examples.
|
|
align_prior_hop_length: Optional int, hop length of audio features.
|
|
If provided alignment prior will be calculated and included in batch output. Must match hop length
|
|
of audio features used for training.
|
|
min_duration: Optional float, if provided audio files in the training manifest shorter than 'min_duration'
|
|
will be ignored.
|
|
max_duration: Optional float, if provided audio files in the training manifest longer than 'max_duration'
|
|
will be ignored.
|
|
volume_norm: Whether to apply volume normalization to loaded audio.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dataset_meta: Dict,
|
|
sample_rate: int,
|
|
text_tokenizer: BaseTokenizer,
|
|
weighted_sampling_steps_per_epoch: Optional[int] = None,
|
|
speaker_path: Optional[Path] = None,
|
|
featurizers: Optional[Dict[str, Featurizer]] = None,
|
|
feature_processors: Optional[Dict[str, FeatureProcessor]] = None,
|
|
align_prior_hop_length: Optional[int] = None,
|
|
min_duration: Optional[float] = None,
|
|
max_duration: Optional[float] = None,
|
|
volume_norm: bool = True,
|
|
):
|
|
super().__init__()
|
|
|
|
self.sample_rate = sample_rate
|
|
self.text_tokenizer = text_tokenizer
|
|
self.weighted_sampling_steps_per_epoch = weighted_sampling_steps_per_epoch
|
|
self.align_prior_hop_length = align_prior_hop_length
|
|
self.include_align_prior = self.align_prior_hop_length is not None
|
|
self.volume_norm = volume_norm
|
|
|
|
if speaker_path:
|
|
self.include_speaker = True
|
|
with open(speaker_path, 'r', encoding="utf-8") as speaker_f:
|
|
speaker_index_map = json.load(speaker_f)
|
|
else:
|
|
self.include_speaker = False
|
|
speaker_index_map = None
|
|
|
|
if featurizers:
|
|
logging.info(f"Found featurizers {featurizers.keys()}")
|
|
self.featurizers = list(featurizers.values())
|
|
else:
|
|
self.featurizers = []
|
|
|
|
if feature_processors:
|
|
logging.info(f"Found featurize processors {feature_processors.keys()}")
|
|
self.feature_processors = list(feature_processors.values())
|
|
else:
|
|
self.feature_processors = []
|
|
|
|
self.data_samples = []
|
|
self.sample_weights = []
|
|
for dataset_name, dataset_info in dataset_meta.items():
|
|
dataset = DatasetMeta(**dataset_info)
|
|
samples, weights = self._preprocess_manifest(
|
|
dataset_name=dataset_name,
|
|
dataset=dataset,
|
|
min_duration=min_duration,
|
|
max_duration=max_duration,
|
|
speaker_index_map=speaker_index_map,
|
|
)
|
|
self.data_samples += samples
|
|
self.sample_weights += weights
|
|
|
|
def get_sampler(self, batch_size: int, world_size: int) -> Optional[torch.utils.data.Sampler]:
|
|
if not self.weighted_sampling_steps_per_epoch:
|
|
return None
|
|
|
|
sampler = get_weighted_sampler(
|
|
sample_weights=self.sample_weights,
|
|
batch_size=batch_size,
|
|
world_size=world_size,
|
|
num_steps=self.weighted_sampling_steps_per_epoch,
|
|
)
|
|
return sampler
|
|
|
|
def _preprocess_manifest(
|
|
self,
|
|
dataset_name: str,
|
|
dataset: DatasetMeta,
|
|
min_duration: float,
|
|
max_duration: float,
|
|
speaker_index_map: Dict[str, int],
|
|
):
|
|
entries = read_manifest(dataset.manifest_path)
|
|
filtered_entries, total_hours, filtered_hours = filter_dataset_by_duration(
|
|
entries=entries, min_duration=min_duration, max_duration=max_duration
|
|
)
|
|
|
|
logging.info(dataset_name)
|
|
logging.info(f"Original # of files: {len(entries)}")
|
|
logging.info(f"Filtered # of files: {len(filtered_entries)}")
|
|
logging.info(f"Original duration: {total_hours:.2f} hours")
|
|
logging.info(f"Filtered duration: {filtered_hours:.2f} hours")
|
|
|
|
samples = []
|
|
sample_weights = []
|
|
for entry in filtered_entries:
|
|
|
|
if "normalized_text" in entry:
|
|
text = entry["normalized_text"]
|
|
else:
|
|
text = entry["text"]
|
|
|
|
if self.include_speaker:
|
|
speaker = entry["speaker"]
|
|
speaker_index = speaker_index_map[speaker]
|
|
else:
|
|
speaker = None
|
|
speaker_index = 0
|
|
|
|
if "language" in entry:
|
|
language = entry.get("language")
|
|
elif dataset.language:
|
|
language = dataset.language
|
|
else:
|
|
language = None
|
|
|
|
sample = DatasetSample(
|
|
dataset_name=dataset_name,
|
|
manifest_entry=entry,
|
|
audio_dir=None if dataset.audio_dir is None else Path(dataset.audio_dir),
|
|
feature_dir=None if dataset.feature_dir is None else Path(dataset.feature_dir),
|
|
text=text,
|
|
speaker=speaker,
|
|
speaker_index=speaker_index,
|
|
language=language,
|
|
tokenizer_names=dataset.tokenizer_names,
|
|
)
|
|
samples.append(sample)
|
|
sample_weights.append(dataset.sample_weight)
|
|
|
|
return samples, sample_weights
|
|
|
|
def __len__(self):
|
|
return len(self.data_samples)
|
|
|
|
def __getitem__(self, index):
|
|
data = self.data_samples[index]
|
|
|
|
audio_array, _, audio_filepath_rel = load_audio(
|
|
manifest_entry=data.manifest_entry,
|
|
audio_dir=data.audio_dir,
|
|
sample_rate=self.sample_rate,
|
|
volume_norm=self.volume_norm,
|
|
)
|
|
audio = torch.tensor(audio_array, dtype=torch.float32)
|
|
audio_len = audio.shape[0]
|
|
|
|
tokens = self.text_tokenizer(data.text)
|
|
tokens = torch.tensor(tokens, dtype=torch.int32)
|
|
text_len = tokens.shape[0]
|
|
|
|
example = {
|
|
"dataset_name": data.dataset_name,
|
|
"audio_filepath": audio_filepath_rel,
|
|
"audio": audio,
|
|
"audio_len": audio_len,
|
|
"tokens": tokens,
|
|
"text_len": text_len,
|
|
}
|
|
|
|
if data.speaker is not None:
|
|
example["speaker"] = data.speaker
|
|
example["speaker_index"] = data.speaker_index
|
|
|
|
if self.include_align_prior:
|
|
spec_len = 1 + librosa.core.samples_to_frames(audio_len, hop_length=self.align_prior_hop_length)
|
|
align_prior = beta_binomial_prior_distribution(phoneme_count=text_len, mel_count=spec_len)
|
|
align_prior = torch.tensor(align_prior, dtype=torch.float32)
|
|
example["align_prior"] = align_prior
|
|
|
|
for featurizer in self.featurizers:
|
|
feature_dict = featurizer.load(
|
|
manifest_entry=data.manifest_entry, audio_dir=data.audio_dir, feature_dir=data.feature_dir
|
|
)
|
|
example.update(feature_dict)
|
|
|
|
for processor in self.feature_processors:
|
|
processor.process(example)
|
|
|
|
return example
|
|
|
|
def collate_fn(self, batch: List[dict]):
|
|
dataset_name_list = []
|
|
audio_filepath_list = []
|
|
audio_list = []
|
|
audio_len_list = []
|
|
token_list = []
|
|
token_len_list = []
|
|
speaker_list = []
|
|
prior_list = []
|
|
|
|
for example in batch:
|
|
dataset_name_list.append(example["dataset_name"])
|
|
audio_filepath_list.append(example["audio_filepath"])
|
|
|
|
audio_list.append(example["audio"])
|
|
audio_len_list.append(example["audio_len"])
|
|
|
|
token_list.append(example["tokens"])
|
|
token_len_list.append(example["text_len"])
|
|
|
|
if self.include_speaker:
|
|
speaker_list.append(example["speaker_index"])
|
|
|
|
if self.include_align_prior:
|
|
prior_list.append(example["align_prior"])
|
|
|
|
batch_audio_len = torch.IntTensor(audio_len_list)
|
|
audio_max_len = int(batch_audio_len.max().item())
|
|
|
|
batch_token_len = torch.IntTensor(token_len_list)
|
|
token_max_len = int(batch_token_len.max().item())
|
|
|
|
batch_audio = stack_tensors(audio_list, max_lens=[audio_max_len])
|
|
batch_tokens = stack_tensors(token_list, max_lens=[token_max_len], pad_value=self.text_tokenizer.pad)
|
|
|
|
batch_dict = {
|
|
"dataset_names": dataset_name_list,
|
|
"audio_filepaths": audio_filepath_list,
|
|
"audio": batch_audio,
|
|
"audio_lens": batch_audio_len,
|
|
"text": batch_tokens,
|
|
"text_lens": batch_token_len,
|
|
}
|
|
|
|
if self.include_speaker:
|
|
batch_dict["speaker_id"] = torch.IntTensor(speaker_list)
|
|
|
|
if self.include_align_prior:
|
|
spec_max_len = max([prior.shape[0] for prior in prior_list])
|
|
text_max_len = max([prior.shape[1] for prior in prior_list])
|
|
batch_dict["align_prior_matrix"] = stack_tensors(
|
|
prior_list,
|
|
max_lens=[text_max_len, spec_max_len],
|
|
)
|
|
|
|
for featurizer in self.featurizers:
|
|
feature_dict = featurizer.collate_fn(batch)
|
|
batch_dict.update(feature_dict)
|
|
|
|
return batch_dict
|
|
|
|
|
|
class MagpieTTSDataset(TextToSpeechDataset):
|
|
"""
|
|
Class for processing and loading text to speech training examples for Magpie-TTS model.
|
|
In addition to the manifest structure for TextToSpeechDataset, we can have the following keys:
|
|
context_audio_filepath, context_audio_duration, target_audio_codes_path, context_audio_codes_path.
|
|
Note: target_audio_codes_path, context_audio_codes_path are absolute paths to the cached audio codes.
|
|
If they are not present in the manifest or if load_cached_codes_if_available=False, then the audio will be loaded
|
|
and codes will be computed on the fly in the model class.
|
|
|
|
Args:
|
|
dataset_meta: Dict of dataset names (string) to dataset metadata.
|
|
sample_rate: Sample rate to load audio as. If the audio is stored at a different sample rate, then it will
|
|
be resampled.
|
|
weighted_sampling_steps_per_epoch: Optional int, If provided, then data will be sampled (with replacement) based on
|
|
the sample weights provided in the dataset metadata. If None, then sample weights will be ignored.
|
|
min_duration: Optional float, if provided audio files in the training manifest shorter than 'min_duration'
|
|
will be ignored.
|
|
max_duration: Optional float, if provided audio files in the training manifest longer than 'max_duration'
|
|
will be ignored.
|
|
volume_norm: Whether to apply volume normalization to loaded audio.
|
|
codec_model_samples_per_frame: Num samples in waveform per codec frame (codec downsample factor).
|
|
bos_id: Text BOS token id.
|
|
eos_id: Text EOS token id.
|
|
num_audio_codebooks: Number of audio codebooks.
|
|
prior_scaling_factor: Scaling factor for the beta binomial prior distribution.
|
|
load_cached_codes_if_available: Whether to load cached audio codes if available *_codes_path keys are available in the manifest.
|
|
dataset_type: Dataset type (train, dev, test).
|
|
tokenizer_config: Config of the tokenzizer used in worker_init_fn to setup the tokenizer (See Magpie-TTS yamls)
|
|
load_16khz_audio: Whether to load 16khz audio for SV model.
|
|
use_text_conditioning_tokenizer: Set True for text context conditioning.
|
|
pad_context_text_to_max_duration: Whether to pad context text to max context audio frames.
|
|
context_duration_min: Minimum duration of context audio in seconds.
|
|
context_duration_max: Maximum duration of context audio in seconds.
|
|
text_context_remapping: Dict defining mapping of multiple text contexts to a single text context.
|
|
text_context_remapping_prob: Probability of remapping the original text context to a remapped text context.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dataset_meta: Dict,
|
|
sample_rate: int,
|
|
weighted_sampling_steps_per_epoch: Optional[int] = None,
|
|
min_duration: Optional[float] = None,
|
|
max_duration: Optional[float] = None,
|
|
volume_norm: bool = True,
|
|
codec_model_samples_per_frame: int = None,
|
|
bos_id: int = None,
|
|
eos_id: int = None,
|
|
num_audio_codebooks: int = None,
|
|
prior_scaling_factor: float = None,
|
|
load_cached_codes_if_available: bool = True,
|
|
dataset_type: str = 'train',
|
|
tokenizer_config=None,
|
|
load_16khz_audio: bool = True,
|
|
use_text_conditioning_tokenizer: bool = False,
|
|
text_conditioning_tokenizer_name: str = None,
|
|
pad_context_text_to_max_duration: bool = False,
|
|
context_duration_min: float = 3.0,
|
|
context_duration_max: float = 10.0,
|
|
text_context_remapping: Dict[str, str] = None,
|
|
text_context_remapping_prob: float = 0.0,
|
|
ignore_phoneme_languages: List[str] = None,
|
|
phoneme_as_text_prob: float = 0.0,
|
|
pronunciation_control_g2p: Dict = None,
|
|
add_language_to_context_text: bool = False,
|
|
default_tokenizer_name: str = "english_phoneme",
|
|
):
|
|
super().__init__(
|
|
dataset_meta=dataset_meta,
|
|
sample_rate=sample_rate,
|
|
text_tokenizer=None,
|
|
weighted_sampling_steps_per_epoch=weighted_sampling_steps_per_epoch,
|
|
speaker_path=None,
|
|
featurizers=None,
|
|
feature_processors=None,
|
|
align_prior_hop_length=None,
|
|
min_duration=min_duration,
|
|
max_duration=max_duration,
|
|
volume_norm=volume_norm,
|
|
)
|
|
self.bos_id = bos_id # TODO @xueyang: this should be removed since no other places used it.
|
|
self.eos_id = eos_id
|
|
self.num_audio_codebooks = num_audio_codebooks
|
|
self.codec_model_samples_per_frame = codec_model_samples_per_frame
|
|
self.include_align_prior = prior_scaling_factor is not None
|
|
self.prior_scaling_factor = prior_scaling_factor
|
|
self.load_cached_codes_if_available = load_cached_codes_if_available
|
|
self.dataset_type = dataset_type
|
|
self.tokenizer_config = tokenizer_config
|
|
self.text_tokenizer = None # Assigned in worker_init_fn in model file
|
|
self.phoneme_tokenizer = None # Assigned in worker_init_fn in model file (if any)
|
|
self.pronunciation_control_g2p = None
|
|
self.load_16khz_audio = load_16khz_audio
|
|
self.use_text_conditioning_tokenizer = use_text_conditioning_tokenizer
|
|
self.text_conditioning_tokenizer_name = text_conditioning_tokenizer_name
|
|
self.pad_context_text_to_max_duration = pad_context_text_to_max_duration
|
|
self.context_duration_min = context_duration_min
|
|
self.context_duration_max = context_duration_max
|
|
self.text_context_remapping = text_context_remapping
|
|
self.text_context_remapping_prob = text_context_remapping_prob
|
|
self.ignore_phoneme_languages = ignore_phoneme_languages or []
|
|
self.phoneme_as_text_prob = phoneme_as_text_prob
|
|
self.pronunciation_control_g2p_config = pronunciation_control_g2p
|
|
self.add_language_to_context_text = add_language_to_context_text
|
|
self.default_tokenizer_name = default_tokenizer_name
|
|
|
|
def get_num_audio_samples_to_slice(self, duration, sample_rate):
|
|
num_codec_frames = int(duration * sample_rate / self.codec_model_samples_per_frame)
|
|
num_audio_samples = num_codec_frames * self.codec_model_samples_per_frame
|
|
return num_audio_samples
|
|
|
|
def __getitem__(self, index):
|
|
data = self.data_samples[index]
|
|
if (
|
|
self.pronunciation_control_g2p is None
|
|
and self.pronunciation_control_g2p_config is not None
|
|
and self.phoneme_as_text_prob > 0.0
|
|
):
|
|
self.pronunciation_control_g2p = setup_pronunciation_control_g2p(self.pronunciation_control_g2p_config)
|
|
|
|
def _sample_context_duration_with_available_limit(available_duration_sec: float) -> float:
|
|
effective_duration_max = min(self.context_duration_max, available_duration_sec)
|
|
effective_duration_max = max(self.context_duration_min, effective_duration_max)
|
|
return random.uniform(self.context_duration_min, effective_duration_max)
|
|
|
|
if data.tokenizer_names is not None:
|
|
# Pick a random tokenizer from the list of tokenizers
|
|
tokenizer_name = random.choice(data.tokenizer_names)
|
|
else:
|
|
tokenizer_name = self.default_tokenizer_name
|
|
|
|
if data.language:
|
|
language = data.language
|
|
else:
|
|
language = 'en'
|
|
|
|
tokens = tokenize_text_with_pronunciation_control(
|
|
text_tokenizer=self.text_tokenizer,
|
|
text_str=data.text,
|
|
language=language,
|
|
tokenizer_name=tokenizer_name,
|
|
dataset_type=self.dataset_type,
|
|
phoneme_as_text_prob=self.phoneme_as_text_prob,
|
|
pronunciation_control_g2p=self.pronunciation_control_g2p,
|
|
)
|
|
tokens = tokens + [self.eos_id] # Not adding BOS id
|
|
tokens = torch.tensor(tokens, dtype=torch.int32)
|
|
text_len = tokens.shape[0]
|
|
|
|
example = {
|
|
"dataset_name": data.dataset_name,
|
|
"tokens": tokens,
|
|
"text_len": text_len,
|
|
}
|
|
|
|
if self.phoneme_tokenizer is not None:
|
|
# Use IPA text for IPABPETokenizer (required), otherwise use regular text
|
|
if isinstance(self.phoneme_tokenizer, IPABPETokenizer):
|
|
if 'ipa' not in data.manifest_entry:
|
|
if self.dataset_type == 'train':
|
|
raise ValueError(
|
|
f"IPABPETokenizer requires 'ipa' field but it is not available in the manifest entry. "
|
|
f"Text: {data.text}"
|
|
)
|
|
else:
|
|
logging.warning(
|
|
f"'ipa' field not found in manifest entry for text: {data.text}. "
|
|
f"Use only predicted phonemes for inference."
|
|
)
|
|
phoneme_text = data.manifest_entry.get('ipa', '')
|
|
if language in self.ignore_phoneme_languages:
|
|
# Ignore phoneme tokenization for this language.
|
|
phoneme_text = ""
|
|
else:
|
|
phoneme_text = data.text
|
|
phoneme_tokens = self.phoneme_tokenizer.encode(phoneme_text)
|
|
phoneme_tokens = (
|
|
[self.phoneme_tokenizer.bos_token_id] + phoneme_tokens + [self.phoneme_tokenizer.eos_token_id]
|
|
)
|
|
phoneme_tokens_len = len(phoneme_tokens)
|
|
example["phoneme_tokens"] = torch.tensor(phoneme_tokens, dtype=torch.int32)
|
|
example["phoneme_tokens_len"] = phoneme_tokens_len
|
|
|
|
if self.load_cached_codes_if_available and 'target_audio_codes_path' in data.manifest_entry:
|
|
audio_codes_path = data.manifest_entry['target_audio_codes_path']
|
|
audio_codes = torch.load(audio_codes_path) # (C, T)
|
|
spec_len = audio_codes.shape[1] + 1 # +1 for EOS
|
|
audio_codes_len = audio_codes.shape[1]
|
|
example['audio_codes'] = audio_codes
|
|
example['audio_codes_len'] = audio_codes_len
|
|
example['audio_filepath'] = audio_codes_path
|
|
if 'audio_filepath' in data.manifest_entry:
|
|
# If audio_filepath is available, then use the actual audio file path.
|
|
example['audio_filepath'] = data.manifest_entry['audio_filepath']
|
|
elif 'audio_filepath' in data.manifest_entry:
|
|
# Only load audio if codes are not available
|
|
audio_array, _, audio_filepath_rel = load_audio(
|
|
manifest_entry=data.manifest_entry,
|
|
audio_dir=data.audio_dir,
|
|
sample_rate=self.sample_rate,
|
|
volume_norm=self.volume_norm,
|
|
)
|
|
audio = torch.tensor(audio_array, dtype=torch.float32)
|
|
# Pad audio to be multiple of downsample factor
|
|
audio = torch.nn.functional.pad(
|
|
audio,
|
|
(0, self.codec_model_samples_per_frame - (audio.shape[0] % self.codec_model_samples_per_frame)),
|
|
value=0,
|
|
)
|
|
audio_len = audio.shape[0]
|
|
example['audio_filepath'] = data.manifest_entry['audio_filepath']
|
|
example['audio'] = audio
|
|
example['audio_len'] = audio_len
|
|
spec_len = int(audio_len / self.codec_model_samples_per_frame) + 1 # +1 for EOS
|
|
|
|
if self.load_cached_codes_if_available and 'context_audio_codes_path' in data.manifest_entry:
|
|
context_audio_codes_path = data.manifest_entry['context_audio_codes_path']
|
|
context_audio_codes = torch.load(context_audio_codes_path) # (8, T)
|
|
_available_context_duration = (
|
|
context_audio_codes.shape[1] * self.codec_model_samples_per_frame / self.sample_rate
|
|
)
|
|
_context_duration_to_slice = _sample_context_duration_with_available_limit(_available_context_duration)
|
|
_num_frames_to_slice = int(
|
|
_context_duration_to_slice * self.sample_rate / self.codec_model_samples_per_frame
|
|
)
|
|
if _num_frames_to_slice < context_audio_codes.shape[1]:
|
|
start_idx = random.randint(0, context_audio_codes.shape[1] - _num_frames_to_slice)
|
|
context_audio_codes = context_audio_codes[:, start_idx : start_idx + _num_frames_to_slice]
|
|
else:
|
|
# Repeaet the audio if it is shorter than the desired duration
|
|
_num_repeats = int(np.ceil(_num_frames_to_slice / context_audio_codes.shape[1]))
|
|
# context_audio_codes is a tensor of shape (num_codebooks, T)
|
|
context_audio_codes_repeated = context_audio_codes.repeat(1, _num_repeats)
|
|
context_audio_codes = context_audio_codes_repeated[:, :_num_frames_to_slice]
|
|
|
|
context_audio_codes_len = context_audio_codes.shape[1]
|
|
example['context_audio_codes'] = context_audio_codes
|
|
example['context_audio_codes_len'] = context_audio_codes_len
|
|
elif 'context_audio_filepath' in data.manifest_entry:
|
|
context_audio_filepath = os.path.join(data.audio_dir, data.manifest_entry['context_audio_filepath'])
|
|
context_duration = data.manifest_entry['context_audio_duration']
|
|
context_audio_array = _read_audio(
|
|
audio_filepath=context_audio_filepath,
|
|
sample_rate=self.sample_rate,
|
|
offset=0,
|
|
duration=context_duration,
|
|
)
|
|
context_audio_array = context_audio_array.samples
|
|
_available_context_duration = len(context_audio_array) / self.sample_rate
|
|
_context_duration_to_slice = _sample_context_duration_with_available_limit(_available_context_duration)
|
|
_num_samples_to_slice = self.get_num_audio_samples_to_slice(_context_duration_to_slice, self.sample_rate)
|
|
if _num_samples_to_slice < len(context_audio_array):
|
|
start_idx = random.randint(0, len(context_audio_array) - _num_samples_to_slice)
|
|
context_audio_array = context_audio_array[start_idx : start_idx + _num_samples_to_slice]
|
|
else:
|
|
# Repeaet the audio if it is shorter than the desired duration
|
|
_num_repeats = int(np.ceil(_num_samples_to_slice / len(context_audio_array)))
|
|
context_audio_array = np.tile(context_audio_array, _num_repeats)
|
|
context_audio_array = context_audio_array[:_num_samples_to_slice]
|
|
context_audio = torch.tensor(context_audio_array, dtype=torch.float32)
|
|
context_audio_len = context_audio.shape[0]
|
|
example['context_audio'] = context_audio
|
|
example['context_audio_len'] = context_audio_len
|
|
else:
|
|
# We always want to have context_audio_codes if available for multi-encoder model. These are ignored for singlencoder model.
|
|
# If context audio is not available, just use a dummy context_audio_codes
|
|
# (Will be used in text context scenario)
|
|
if self.load_cached_codes_if_available:
|
|
context_audio_codes = torch.zeros([self.num_audio_codebooks, 0], dtype=torch.int32)
|
|
context_audio_codes_len = 0
|
|
example['context_audio_codes'] = context_audio_codes
|
|
example['context_audio_codes_len'] = context_audio_codes_len
|
|
else:
|
|
# @shehzeenh: Added this condition so that a batch does not have a mix of context_audio and context_audio_codes
|
|
# @blisc: Added a +1. If we send in exactly 882 samples, then a conv layer complains about padding.
|
|
# Adding 883 works. This occurs when we use text context during inference.
|
|
context_audio = torch.zeros(self.codec_model_samples_per_frame + 1, dtype=torch.float32)
|
|
context_audio_len = context_audio.shape[0]
|
|
example['context_audio'] = context_audio
|
|
example['context_audio_len'] = context_audio_len
|
|
|
|
# 16kHz audio is used for SV model
|
|
if self.load_16khz_audio:
|
|
if 'context_audio_filepath' in data.manifest_entry:
|
|
# If context_audio_filepath is available, then use that for 16khz audio for SV model
|
|
context_audio_filepath = os.path.join(data.audio_dir, data.manifest_entry['context_audio_filepath'])
|
|
context_duration = data.manifest_entry['context_audio_duration']
|
|
audio_array_16khz = _read_audio(
|
|
audio_filepath=context_audio_filepath, sample_rate=16000, offset=0, duration=context_duration
|
|
).samples
|
|
else:
|
|
# Otherwise, load the target audio file.
|
|
audio_array_16khz, _, _ = load_audio(
|
|
manifest_entry=data.manifest_entry,
|
|
audio_dir=data.audio_dir,
|
|
sample_rate=16000,
|
|
volume_norm=self.volume_norm,
|
|
)
|
|
_available_context_duration = len(audio_array_16khz) / 16000
|
|
_context_duration_to_slice = _sample_context_duration_with_available_limit(_available_context_duration)
|
|
_num_samples_to_slice = int(_context_duration_to_slice * 16000)
|
|
if _num_samples_to_slice < len(audio_array_16khz):
|
|
start_idx = random.randint(0, len(audio_array_16khz) - _num_samples_to_slice)
|
|
audio_array_16khz = audio_array_16khz[start_idx : start_idx + _num_samples_to_slice]
|
|
audio_16khz = torch.tensor(audio_array_16khz, dtype=torch.float32)
|
|
audio_len_16khz = audio_16khz.shape[0]
|
|
example['audio_16khz'] = audio_16khz
|
|
example['audio_len_16khz'] = audio_len_16khz
|
|
|
|
if self.use_text_conditioning_tokenizer:
|
|
if 'context_text' in data.manifest_entry:
|
|
context_text = data.manifest_entry['context_text']
|
|
if self.text_context_remapping is not None and context_text in self.text_context_remapping:
|
|
if self.dataset_type == 'train' and random.random() < self.text_context_remapping_prob:
|
|
# Only remap during training. Give the exact text context during inference.
|
|
context_text = self.text_context_remapping[context_text]
|
|
context_tokens = self.text_tokenizer.encode(context_text, self.text_conditioning_tokenizer_name)
|
|
example['has_text_context'] = True
|
|
else:
|
|
if self.add_language_to_context_text:
|
|
context_text = f"[{language.upper()}]"
|
|
else:
|
|
context_text = "[NO TEXT CONTEXT]"
|
|
context_tokens = self.text_tokenizer.encode(context_text, self.text_conditioning_tokenizer_name)
|
|
example['has_text_context'] = False
|
|
if self.pad_context_text_to_max_duration:
|
|
_required_len = (
|
|
int(self.context_duration_max * self.sample_rate / self.codec_model_samples_per_frame) + 2
|
|
) # +2 for BOS and EOS
|
|
if len(context_tokens) < _required_len:
|
|
_pad_id = self.text_tokenizer.tokenizer_pad_ids[self.text_conditioning_tokenizer_name]
|
|
context_tokens += [_pad_id] * (_required_len - len(context_tokens))
|
|
else:
|
|
context_tokens = context_tokens[:_required_len]
|
|
|
|
context_tokens = torch.tensor(context_tokens, dtype=torch.int32)
|
|
context_text_len = context_tokens.shape[0]
|
|
example['context_text_tokens'] = context_tokens
|
|
example['context_text_len'] = context_text_len
|
|
|
|
if self.include_align_prior:
|
|
align_prior = beta_binomial_prior_distribution(
|
|
phoneme_count=text_len, mel_count=spec_len, scaling_factor=self.prior_scaling_factor
|
|
)
|
|
align_prior = torch.tensor(align_prior, dtype=torch.float32)
|
|
example["align_prior"] = align_prior
|
|
|
|
if "original_text" in data.manifest_entry:
|
|
# Raw Text is used as the GT for CER/WER computation in DPO pref data generation
|
|
# and GRPO reward setup. For manifests in which the 'text' field is phonemized,
|
|
# we use the 'original_text' field as the raw text. Otherwise, we use the regular text field.
|
|
example['raw_text'] = data.manifest_entry['original_text']
|
|
else:
|
|
example['raw_text'] = data.text
|
|
|
|
example['language'] = language
|
|
|
|
if "reward" in data.manifest_entry:
|
|
example["reward"] = data.manifest_entry["reward"]
|
|
|
|
# Get speaker index if available (for models with baked context embeddings)
|
|
if 'speaker_index' in data.manifest_entry:
|
|
example['speaker_index'] = data.manifest_entry['speaker_index']
|
|
|
|
return example
|
|
|
|
def collate_fn(self, batch: List[dict]):
|
|
dataset_name_list = []
|
|
audio_filepath_list = []
|
|
audio_list = []
|
|
audio_len_list = []
|
|
audio_list_16khz = []
|
|
audio_len_list_16khz = []
|
|
token_list = []
|
|
token_len_list = []
|
|
prior_list = []
|
|
audio_codes_list = []
|
|
audio_codes_len_list = []
|
|
context_audio_list = []
|
|
context_audio_len_list = []
|
|
context_audio_codes_list = []
|
|
context_audio_codes_len_list = []
|
|
context_text_tokens_list = []
|
|
context_text_tokens_len_list = []
|
|
context_has_text_context_list = []
|
|
reward_list = []
|
|
raw_text_list = []
|
|
language_list = []
|
|
speaker_indices_list = []
|
|
phoneme_tokens_list = []
|
|
phoneme_tokens_len_list = []
|
|
for example in batch:
|
|
dataset_name_list.append(example["dataset_name"])
|
|
raw_text_list.append(example["raw_text"])
|
|
language_list.append(example["language"])
|
|
|
|
token_list.append(example["tokens"])
|
|
token_len_list.append(example["text_len"])
|
|
|
|
if 'audio_filepath' in example:
|
|
audio_filepath_list.append(example["audio_filepath"])
|
|
if 'phoneme_tokens' in example:
|
|
phoneme_tokens_list.append(example["phoneme_tokens"])
|
|
phoneme_tokens_len_list.append(example["phoneme_tokens_len"])
|
|
|
|
if 'audio' in example:
|
|
audio_list.append(example["audio"])
|
|
audio_len_list.append(example["audio_len"])
|
|
|
|
if 'audio_16khz' in example:
|
|
audio_list_16khz.append(example["audio_16khz"])
|
|
audio_len_list_16khz.append(example["audio_len_16khz"])
|
|
|
|
if 'audio_codes' in example:
|
|
audio_codes_list.append(example['audio_codes'])
|
|
audio_codes_len_list.append(example['audio_codes_len'])
|
|
|
|
if 'context_audio' in example:
|
|
context_audio_list.append(example['context_audio'])
|
|
context_audio_len_list.append(example['context_audio_len'])
|
|
|
|
if 'context_audio_codes' in example:
|
|
context_audio_codes_list.append(example['context_audio_codes'])
|
|
context_audio_codes_len_list.append(example['context_audio_codes_len'])
|
|
|
|
if 'context_text_tokens' in example:
|
|
context_text_tokens_list.append(example['context_text_tokens'])
|
|
context_text_tokens_len_list.append(example['context_text_len'])
|
|
context_has_text_context_list.append(example['has_text_context'])
|
|
|
|
if 'reward' in example:
|
|
reward_list.append(example['reward'])
|
|
|
|
if 'speaker_index' in example:
|
|
speaker_indices_list.append(example['speaker_index'])
|
|
|
|
if self.include_align_prior:
|
|
prior_list.append(example["align_prior"])
|
|
|
|
batch_token_len = torch.IntTensor(token_len_list)
|
|
token_max_len = int(batch_token_len.max().item())
|
|
batch_tokens = stack_tensors(token_list, max_lens=[token_max_len], pad_value=self.text_tokenizer.pad)
|
|
|
|
batch_dict = {
|
|
"dataset_names": dataset_name_list,
|
|
"raw_texts": raw_text_list,
|
|
"languages": language_list,
|
|
"audio_filepaths": audio_filepath_list,
|
|
"text": batch_tokens,
|
|
"text_lens": batch_token_len,
|
|
}
|
|
|
|
if len(audio_list) > 0:
|
|
batch_audio_len = torch.IntTensor(audio_len_list)
|
|
audio_max_len = int(batch_audio_len.max().item())
|
|
batch_audio = stack_tensors(audio_list, max_lens=[audio_max_len])
|
|
batch_dict['audio'] = batch_audio
|
|
batch_dict['audio_lens'] = batch_audio_len
|
|
|
|
if len(audio_list_16khz) > 0:
|
|
batch_audio_len_16khz = torch.IntTensor(audio_len_list_16khz)
|
|
audio_max_len_16khz = int(batch_audio_len_16khz.max().item())
|
|
batch_audio_16khz = stack_tensors(audio_list_16khz, max_lens=[audio_max_len_16khz])
|
|
batch_dict['audio_16khz'] = batch_audio_16khz
|
|
batch_dict['audio_lens_16khz'] = batch_audio_len_16khz
|
|
|
|
if len(audio_codes_list) > 0:
|
|
batch_audio_codes_len = torch.IntTensor(audio_codes_len_list)
|
|
audio_codes_max_len = int(batch_audio_codes_len.max().item())
|
|
batch_audio_codes = stack_tensors(audio_codes_list, max_lens=[audio_codes_max_len])
|
|
batch_dict['audio_codes'] = batch_audio_codes
|
|
batch_dict['audio_codes_lens'] = batch_audio_codes_len
|
|
|
|
if len(phoneme_tokens_list) > 0:
|
|
batch_phoneme_tokens_len = torch.IntTensor(phoneme_tokens_len_list)
|
|
phoneme_tokens_max_len = int(batch_phoneme_tokens_len.max().item())
|
|
batch_phoneme_tokens = stack_tensors(
|
|
phoneme_tokens_list, max_lens=[phoneme_tokens_max_len], pad_value=self.phoneme_tokenizer.pad
|
|
)
|
|
batch_dict['phoneme_tokens'] = batch_phoneme_tokens
|
|
batch_dict['phoneme_tokens_lens'] = batch_phoneme_tokens_len
|
|
|
|
if len(context_audio_list) > 0:
|
|
batch_context_audio_len = torch.IntTensor(context_audio_len_list)
|
|
context_audio_max_len = int(batch_context_audio_len.max().item())
|
|
batch_context_audio = stack_tensors(context_audio_list, max_lens=[context_audio_max_len])
|
|
batch_dict['context_audio'] = batch_context_audio
|
|
batch_dict['context_audio_lens'] = batch_context_audio_len
|
|
|
|
if len(context_audio_codes_list) > 0:
|
|
batch_context_audio_codes_len = torch.IntTensor(context_audio_codes_len_list)
|
|
context_audio_codes_max_len = int(batch_context_audio_codes_len.max().item())
|
|
# TODO @xueyang: verify if batch_context_audio_codes are integer.
|
|
batch_context_audio_codes = stack_tensors(context_audio_codes_list, max_lens=[context_audio_codes_max_len])
|
|
batch_dict['context_audio_codes'] = batch_context_audio_codes
|
|
batch_dict['context_audio_codes_lens'] = batch_context_audio_codes_len
|
|
|
|
if self.use_text_conditioning_tokenizer:
|
|
batch_context_text_tokens_len = torch.IntTensor(context_text_tokens_len_list)
|
|
context_text_tokens_max_len = int(batch_context_text_tokens_len.max().item())
|
|
# TODO @xueyang: potential bugs if self.tokenizer.pad is not 0.0. verify if batch_context_text_tokens are integer.
|
|
batch_context_text_tokens = stack_tensors(context_text_tokens_list, max_lens=[context_text_tokens_max_len])
|
|
batch_dict['context_text_tokens'] = batch_context_text_tokens
|
|
batch_dict['context_text_tokens_lens'] = batch_context_text_tokens_len
|
|
batch_dict['has_text_context'] = torch.BoolTensor(context_has_text_context_list)
|
|
|
|
if self.include_align_prior:
|
|
spec_max_len = max([prior.shape[0] for prior in prior_list])
|
|
text_max_len = max([prior.shape[1] for prior in prior_list])
|
|
batch_dict["align_prior_matrix"] = stack_tensors(
|
|
prior_list,
|
|
max_lens=[text_max_len, spec_max_len],
|
|
)
|
|
|
|
if len(reward_list) > 0:
|
|
batch_dict['rewards'] = torch.FloatTensor(reward_list)
|
|
|
|
if len(speaker_indices_list) > 0:
|
|
batch_dict['speaker_indices'] = torch.tensor(speaker_indices_list, dtype=torch.int64)
|
|
|
|
# Assert no more than one of audio or audio_codes in the batch
|
|
if 'audio' in batch_dict:
|
|
assert 'audio_codes' not in batch_dict
|
|
|
|
# Assert no more than one of context_audio or context_audio_codes in the batch
|
|
if 'context_audio' in batch_dict:
|
|
assert 'context_audio_codes' not in batch_dict
|
|
|
|
return batch_dict
|
|
|
|
|
|
class MagpieTTSDatasetDPO(MagpieTTSDataset):
|
|
"""
|
|
This class is meant to be used with the DPO model. To generate manifests for this dataset, please use
|
|
- scripts/magpietts/dpo/create_text_contextpairs.py
|
|
- scripts/magpietts/dpo/create_preference_pairs.py
|
|
in sequence to generate samples and create preference pairs.
|
|
"""
|
|
|
|
def __len__(self):
|
|
return len(self.data_samples) // 2
|
|
|
|
def __getitem__(self, index):
|
|
chosen_example = super().__getitem__(index * 2)
|
|
rejected_example = super().__getitem__(index * 2 + 1)
|
|
assert chosen_example['reward'] == 1.0
|
|
assert rejected_example['reward'] < 1.0
|
|
return {"chosen": chosen_example, "rejected": rejected_example}
|
|
|
|
def collate_fn(self, batch: List[dict]):
|
|
chosen_batch = [example['chosen'] for example in batch]
|
|
rejected_batch = [example['rejected'] for example in batch]
|
|
chosen_collated = super().collate_fn(chosen_batch)
|
|
rejected_collated = super().collate_fn(rejected_batch)
|
|
return {"chosen": chosen_collated, "rejected": rejected_collated}
|
|
|
|
|
|
class ChunkedTTSInferenceDataset(MagpieTTSDataset):
|
|
"""
|
|
Unified dataset for TTS inference with automatic text chunking.
|
|
|
|
Inherits from MagpieTTSDataset to reuse context audio loading, text conditioning,
|
|
and other preprocessing logic. Uses language-aware chunking to automatically
|
|
decide whether to split text into sentences:
|
|
- Short text (below language threshold): returns single chunk
|
|
- Long text (above language threshold): returns multiple sentence chunks (multi-chunk)
|
|
|
|
Both language (for threshold) and tokenizer are determined per-sample:
|
|
- Language from manifest's 'language' field
|
|
- Tokenizer from sample's tokenizer_names or mapped from language
|
|
|
|
Args:
|
|
dataset_meta: Dataset metadata dictionary (same format as MagpieTTSDataset).
|
|
sample_rate: Audio sample rate.
|
|
codec_model_samples_per_frame: Samples per codec frame.
|
|
eos_id: End-of-sequence token ID.
|
|
num_audio_codebooks: Number of audio codebooks.
|
|
context_duration_min: Minimum context duration in seconds.
|
|
context_duration_max: Maximum context duration in seconds.
|
|
use_text_conditioning_tokenizer: Whether model uses text conditioning encoder.
|
|
text_conditioning_tokenizer_name: Name of text conditioning tokenizer.
|
|
pad_context_text_to_max_duration: Whether to pad context text.
|
|
load_16khz_audio: Whether to load 16kHz audio for SV model.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dataset_meta: Dict[str, Any],
|
|
sample_rate: int,
|
|
codec_model_samples_per_frame: int,
|
|
eos_id: int,
|
|
num_audio_codebooks: int,
|
|
context_duration_min: float = 3.0,
|
|
context_duration_max: float = 10.0,
|
|
use_text_conditioning_tokenizer: bool = False,
|
|
text_conditioning_tokenizer_name: str = None,
|
|
pad_context_text_to_max_duration: bool = False,
|
|
load_16khz_audio: bool = False,
|
|
**kwargs,
|
|
):
|
|
# Initialize parent - handles manifest reading and context audio loading
|
|
super().__init__(
|
|
dataset_meta=dataset_meta,
|
|
sample_rate=sample_rate,
|
|
codec_model_samples_per_frame=codec_model_samples_per_frame,
|
|
eos_id=eos_id,
|
|
num_audio_codebooks=num_audio_codebooks,
|
|
context_duration_min=context_duration_min,
|
|
context_duration_max=context_duration_max,
|
|
use_text_conditioning_tokenizer=use_text_conditioning_tokenizer,
|
|
text_conditioning_tokenizer_name=text_conditioning_tokenizer_name,
|
|
pad_context_text_to_max_duration=pad_context_text_to_max_duration,
|
|
load_16khz_audio=load_16khz_audio,
|
|
load_cached_codes_if_available=True, # Prefer codes for inference
|
|
dataset_type='test',
|
|
**kwargs,
|
|
)
|
|
|
|
def _get_tokenizer_name(self, data, language: str = "en") -> str:
|
|
"""Resolve the tokenizer name for a single sample.
|
|
|
|
Resolution order:
|
|
1. ``data.tokenizer_names[0]`` — explicit per-sample list from the
|
|
dataset config (first entry chosen for determinism).
|
|
2. Language-based lookup via ``get_tokenizer_for_language(language,
|
|
available)`` using the keys registered in ``self.text_tokenizer``.
|
|
3. Hard-coded fallback ``"english_phoneme"`` when no tokenizer object
|
|
is available at all.
|
|
|
|
Args:
|
|
data: DatasetSample with optional tokenizer_names field.
|
|
language: Language code from the manifest entry (e.g. ``"en"``).
|
|
|
|
Returns:
|
|
Tokenizer name to use for encoding.
|
|
"""
|
|
# First try sample's tokenizer_names (from dataset config)
|
|
if data.tokenizer_names is not None:
|
|
return data.tokenizer_names[0] # Use first (deterministic for inference)
|
|
|
|
# Fall back to centralized language-based mapping
|
|
if self.text_tokenizer is not None:
|
|
available = list(self.text_tokenizer.tokenizers.keys())
|
|
return get_tokenizer_for_language(language, available)
|
|
|
|
return "english_phoneme"
|
|
|
|
def __getitem__(self, idx: int) -> Dict[str, Any]:
|
|
"""
|
|
Add automatic text chunking on top of parent's __getitem__.
|
|
|
|
Uses language-aware chunking that automatically decides whether to split:
|
|
- Short text (below threshold): returns as single chunk
|
|
- Long text (above threshold): returns as multiple sentence chunks
|
|
|
|
Both tokenizer and chunking threshold are determined per-sample based on
|
|
language and tokenizer configuration.
|
|
|
|
Returns:
|
|
Dictionary containing all parent fields plus:
|
|
- idx: Sample index
|
|
- chunked_tokens: List of tokenized text chunks (1 for short, N for long)
|
|
- chunked_tokens_len: List of token lengths
|
|
"""
|
|
# Get data sample for text and tokenizer info
|
|
data = self.data_samples[idx]
|
|
text = data.text
|
|
|
|
# Call parent to get ALL the context audio, text conditioning, etc.
|
|
example = super().__getitem__(idx)
|
|
|
|
# Get language and tokenizer per-sample
|
|
language = example["language"]
|
|
tokenizer_name = self._get_tokenizer_name(data, language)
|
|
|
|
# Unified chunking: automatically decides whether to split based on language threshold
|
|
chunked_tokens, chunked_tokens_len, _ = chunk_text_for_inference(
|
|
text=text,
|
|
language=language,
|
|
tokenizer_name=tokenizer_name,
|
|
text_tokenizer=self.text_tokenizer,
|
|
eos_token_id=self.eos_id,
|
|
)
|
|
|
|
# Handle empty text edge case
|
|
if not chunked_tokens:
|
|
chunked_tokens = [torch.tensor([self.eos_id], dtype=torch.int32)]
|
|
chunked_tokens_len = [1]
|
|
|
|
# Add chunking-related fields
|
|
example['idx'] = idx
|
|
example['chunked_tokens'] = chunked_tokens
|
|
example['chunked_tokens_len'] = chunked_tokens_len
|
|
|
|
return example
|
|
|
|
def collate_fn(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
"""
|
|
Collate function for batching unified inference samples.
|
|
|
|
Calls parent's collate_fn to handle context audio, text conditioning, etc.,
|
|
then adds chunking-related fields (chunked_tokens).
|
|
|
|
Handles mixed batches where samples have different numbers of chunks by
|
|
padding shorter samples with EOS tokens.
|
|
"""
|
|
# Call parent's collate_fn to handle all standard fields
|
|
batch_dict = super().collate_fn(batch)
|
|
|
|
# Add chunking-related fields
|
|
indices = []
|
|
chunked_tokens_list = []
|
|
chunked_tokens_lens_list = []
|
|
|
|
# Find max number of chunks across batch
|
|
max_num_chunks = max(len(sample['chunked_tokens']) for sample in batch)
|
|
|
|
for sample in batch:
|
|
indices.append(sample['idx'])
|
|
|
|
# Pad chunked tokens to max_num_chunks with single EOS token
|
|
num_padding = max_num_chunks - len(sample['chunked_tokens'])
|
|
padded_tokens = sample['chunked_tokens'] + [
|
|
torch.tensor([self.eos_id], dtype=torch.int32) for _ in range(num_padding)
|
|
]
|
|
padded_lens = sample['chunked_tokens_len'] + [1] * num_padding
|
|
|
|
chunked_tokens_list.append(padded_tokens)
|
|
chunked_tokens_lens_list.append(padded_lens)
|
|
|
|
# Add chunking-related fields to batch_dict
|
|
batch_dict['idx'] = indices
|
|
batch_dict['chunked_tokens'] = chunked_tokens_list
|
|
batch_dict['chunked_tokens_lens'] = chunked_tokens_lens_list
|
|
|
|
return batch_dict
|