chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,93 @@
|
||||
import os.path as op
|
||||
from typing import BinaryIO, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def get_waveform(
|
||||
path_or_fp: Union[str, BinaryIO], normalization=True
|
||||
) -> Tuple[np.ndarray, int]:
|
||||
"""Get the waveform and sample rate of a 16-bit mono-channel WAV or FLAC.
|
||||
|
||||
Args:
|
||||
path_or_fp (str or BinaryIO): the path or file-like object
|
||||
normalization (bool): Normalize values to [-1, 1] (Default: True)
|
||||
"""
|
||||
if isinstance(path_or_fp, str):
|
||||
ext = op.splitext(op.basename(path_or_fp))[1]
|
||||
if ext not in {".flac", ".wav"}:
|
||||
raise ValueError(f"Unsupported audio format: {ext}")
|
||||
|
||||
try:
|
||||
import soundfile as sf
|
||||
except ImportError:
|
||||
raise ImportError("Please install soundfile to load WAV/FLAC file")
|
||||
|
||||
waveform, sample_rate = sf.read(path_or_fp, dtype="float32")
|
||||
if not normalization:
|
||||
waveform *= 2 ** 15 # denormalized to 16-bit signed integers
|
||||
return waveform, sample_rate
|
||||
|
||||
|
||||
def _get_kaldi_fbank(waveform, sample_rate, n_bins=80) -> Optional[np.ndarray]:
|
||||
"""Get mel-filter bank features via PyKaldi."""
|
||||
try:
|
||||
from kaldi.feat.mel import MelBanksOptions
|
||||
from kaldi.feat.fbank import FbankOptions, Fbank
|
||||
from kaldi.feat.window import FrameExtractionOptions
|
||||
from kaldi.matrix import Vector
|
||||
|
||||
mel_opts = MelBanksOptions()
|
||||
mel_opts.num_bins = n_bins
|
||||
frame_opts = FrameExtractionOptions()
|
||||
frame_opts.samp_freq = sample_rate
|
||||
opts = FbankOptions()
|
||||
opts.mel_opts = mel_opts
|
||||
opts.frame_opts = frame_opts
|
||||
fbank = Fbank(opts=opts)
|
||||
features = fbank.compute(Vector(waveform), 1.0).numpy()
|
||||
return features
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
|
||||
def _get_torchaudio_fbank(waveform, sample_rate, n_bins=80) -> Optional[np.ndarray]:
|
||||
"""Get mel-filter bank features via TorchAudio."""
|
||||
try:
|
||||
import torch
|
||||
import torchaudio.compliance.kaldi as ta_kaldi
|
||||
import torchaudio.sox_effects as ta_sox
|
||||
|
||||
waveform = torch.from_numpy(waveform)
|
||||
if len(waveform.shape) == 1:
|
||||
# Mono channel: D -> 1 x D
|
||||
waveform = waveform.unsqueeze(0)
|
||||
else:
|
||||
# Merge multiple channels to one: C x D -> 1 x D
|
||||
waveform, _ = ta_sox.apply_effects_tensor(waveform, sample_rate, ['channels', '1'])
|
||||
|
||||
features = ta_kaldi.fbank(
|
||||
waveform, num_mel_bins=n_bins, sample_frequency=sample_rate
|
||||
)
|
||||
return features.numpy()
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
|
||||
def get_fbank(path_or_fp: Union[str, BinaryIO], n_bins=80) -> np.ndarray:
|
||||
"""Get mel-filter bank features via PyKaldi or TorchAudio. Prefer PyKaldi
|
||||
(faster CPP implementation) to TorchAudio (Python implementation). Note that
|
||||
Kaldi/TorchAudio requires 16-bit signed integers as inputs and hence the
|
||||
waveform should not be normalized."""
|
||||
sound, sample_rate = get_waveform(path_or_fp, normalization=False)
|
||||
|
||||
features = _get_kaldi_fbank(sound, sample_rate, n_bins)
|
||||
if features is None:
|
||||
features = _get_torchaudio_fbank(sound, sample_rate, n_bins)
|
||||
if features is None:
|
||||
raise ImportError(
|
||||
"Please install pyKaldi or torchaudio to enable "
|
||||
"online filterbank feature extraction"
|
||||
)
|
||||
|
||||
return features
|
||||
@@ -0,0 +1,82 @@
|
||||
import importlib
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, Optional
|
||||
|
||||
|
||||
class AudioFeatureTransform(ABC):
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def from_config_dict(cls, config: Optional[Dict] = None):
|
||||
pass
|
||||
|
||||
|
||||
AUDIO_FEATURE_TRANSFORM_REGISTRY = {}
|
||||
AUDIO_FEATURE_TRANSFORM_CLASS_NAMES = set()
|
||||
|
||||
|
||||
def register_audio_feature_transform(name):
|
||||
def register_audio_feature_transform_cls(cls):
|
||||
if name in AUDIO_FEATURE_TRANSFORM_REGISTRY:
|
||||
raise ValueError(f"Cannot register duplicate transform ({name})")
|
||||
if not issubclass(cls, AudioFeatureTransform):
|
||||
raise ValueError(
|
||||
f"Transform ({name}: {cls.__name__}) must extend "
|
||||
"AudioFeatureTransform"
|
||||
)
|
||||
if cls.__name__ in AUDIO_FEATURE_TRANSFORM_CLASS_NAMES:
|
||||
raise ValueError(
|
||||
f"Cannot register audio feature transform with duplicate "
|
||||
f"class name ({cls.__name__})"
|
||||
)
|
||||
AUDIO_FEATURE_TRANSFORM_REGISTRY[name] = cls
|
||||
AUDIO_FEATURE_TRANSFORM_CLASS_NAMES.add(cls.__name__)
|
||||
return cls
|
||||
|
||||
return register_audio_feature_transform_cls
|
||||
|
||||
|
||||
def get_audio_feature_transform(name):
|
||||
return AUDIO_FEATURE_TRANSFORM_REGISTRY[name]
|
||||
|
||||
|
||||
transforms_dir = os.path.dirname(__file__)
|
||||
for file in os.listdir(transforms_dir):
|
||||
path = os.path.join(transforms_dir, file)
|
||||
if (
|
||||
not file.startswith("_")
|
||||
and not file.startswith(".")
|
||||
and (file.endswith(".py") or os.path.isdir(path))
|
||||
):
|
||||
name = file[: file.find(".py")] if file.endswith(".py") else file
|
||||
importlib.import_module("fairseq.data.audio.feature_transforms." + name)
|
||||
|
||||
|
||||
class CompositeAudioFeatureTransform(AudioFeatureTransform):
|
||||
@classmethod
|
||||
def from_config_dict(cls, config=None):
|
||||
_config = {} if config is None else config
|
||||
_transforms = _config.get("transforms")
|
||||
if _transforms is None:
|
||||
return None
|
||||
transforms = [
|
||||
get_audio_feature_transform(_t).from_config_dict(_config.get(_t))
|
||||
for _t in _transforms
|
||||
]
|
||||
return CompositeAudioFeatureTransform(transforms)
|
||||
|
||||
def __init__(self, transforms):
|
||||
self.transforms = [t for t in transforms if t is not None]
|
||||
|
||||
def __call__(self, x):
|
||||
for t in self.transforms:
|
||||
x = t(x)
|
||||
return x
|
||||
|
||||
def __repr__(self):
|
||||
format_string = (
|
||||
[self.__class__.__name__ + "("]
|
||||
+ [f" {t.__repr__()}" for t in self.transforms]
|
||||
+ [")"]
|
||||
)
|
||||
return "\n".join(format_string)
|
||||
@@ -0,0 +1,29 @@
|
||||
import numpy as np
|
||||
from fairseq.data.audio.feature_transforms import (
|
||||
AudioFeatureTransform,
|
||||
register_audio_feature_transform,
|
||||
)
|
||||
|
||||
|
||||
@register_audio_feature_transform("global_cmvn")
|
||||
class GlobalCMVN(AudioFeatureTransform):
|
||||
"""Global CMVN (cepstral mean and variance normalization). The global mean
|
||||
and variance need to be pre-computed and stored in NumPy format (.npz)."""
|
||||
|
||||
@classmethod
|
||||
def from_config_dict(cls, config=None):
|
||||
_config = {} if config is None else config
|
||||
return GlobalCMVN(_config.get("stats_npz_path"))
|
||||
|
||||
def __init__(self, stats_npz_path):
|
||||
self.stats_npz_path = stats_npz_path
|
||||
stats = np.load(stats_npz_path)
|
||||
self.mean, self.std = stats["mean"], stats["std"]
|
||||
|
||||
def __repr__(self):
|
||||
return self.__class__.__name__ + f'(stats_npz_path="{self.stats_npz_path}")'
|
||||
|
||||
def __call__(self, x):
|
||||
x = np.subtract(x, self.mean)
|
||||
x = np.divide(x, self.std)
|
||||
return x
|
||||
@@ -0,0 +1,131 @@
|
||||
import math
|
||||
import numbers
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
from fairseq.data.audio.feature_transforms import (
|
||||
AudioFeatureTransform,
|
||||
register_audio_feature_transform,
|
||||
)
|
||||
|
||||
|
||||
@register_audio_feature_transform("specaugment")
|
||||
class SpecAugmentTransform(AudioFeatureTransform):
|
||||
"""SpecAugment (https://arxiv.org/abs/1904.08779)"""
|
||||
|
||||
@classmethod
|
||||
def from_config_dict(cls, config=None):
|
||||
_config = {} if config is None else config
|
||||
return SpecAugmentTransform(
|
||||
_config.get("time_warp_W", 0),
|
||||
_config.get("freq_mask_N", 0),
|
||||
_config.get("freq_mask_F", 0),
|
||||
_config.get("time_mask_N", 0),
|
||||
_config.get("time_mask_T", 0),
|
||||
_config.get("time_mask_p", 0.0),
|
||||
_config.get("mask_value", None),
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
time_warp_w: int = 0,
|
||||
freq_mask_n: int = 0,
|
||||
freq_mask_f: int = 0,
|
||||
time_mask_n: int = 0,
|
||||
time_mask_t: int = 0,
|
||||
time_mask_p: float = 0.0,
|
||||
mask_value: Optional[float] = 0.0,
|
||||
):
|
||||
# Sanity checks
|
||||
assert mask_value is None or isinstance(
|
||||
mask_value, numbers.Number
|
||||
), f"mask_value (type: {type(mask_value)}) must be None or a number"
|
||||
if freq_mask_n > 0:
|
||||
assert freq_mask_f > 0, (
|
||||
f"freq_mask_F ({freq_mask_f}) "
|
||||
f"must be larger than 0 when doing freq masking."
|
||||
)
|
||||
if time_mask_n > 0:
|
||||
assert time_mask_t > 0, (
|
||||
f"time_mask_T ({time_mask_t}) must be larger than 0 when "
|
||||
f"doing time masking."
|
||||
)
|
||||
|
||||
self.time_warp_w = time_warp_w
|
||||
self.freq_mask_n = freq_mask_n
|
||||
self.freq_mask_f = freq_mask_f
|
||||
self.time_mask_n = time_mask_n
|
||||
self.time_mask_t = time_mask_t
|
||||
self.time_mask_p = time_mask_p
|
||||
self.mask_value = mask_value
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
self.__class__.__name__
|
||||
+ "("
|
||||
+ ", ".join(
|
||||
[
|
||||
f"time_warp_w={self.time_warp_w}",
|
||||
f"freq_mask_n={self.freq_mask_n}",
|
||||
f"freq_mask_f={self.freq_mask_f}",
|
||||
f"time_mask_n={self.time_mask_n}",
|
||||
f"time_mask_t={self.time_mask_t}",
|
||||
f"time_mask_p={self.time_mask_p}",
|
||||
]
|
||||
)
|
||||
+ ")"
|
||||
)
|
||||
|
||||
def __call__(self, spectrogram):
|
||||
assert len(spectrogram.shape) == 2, "spectrogram must be a 2-D tensor."
|
||||
|
||||
distorted = spectrogram.copy() # make a copy of input spectrogram.
|
||||
num_frames = spectrogram.shape[0] # or 'tau' in the paper.
|
||||
num_freqs = spectrogram.shape[1] # or 'miu' in the paper.
|
||||
mask_value = self.mask_value
|
||||
|
||||
if mask_value is None: # if no value was specified, use local mean.
|
||||
mask_value = spectrogram.mean()
|
||||
|
||||
if num_frames == 0:
|
||||
return spectrogram
|
||||
|
||||
if num_freqs < self.freq_mask_f:
|
||||
return spectrogram
|
||||
|
||||
if self.time_warp_w > 0:
|
||||
if 2 * self.time_warp_w < num_frames:
|
||||
import cv2
|
||||
|
||||
w0 = np.random.randint(self.time_warp_w, num_frames - self.time_warp_w)
|
||||
w = np.random.randint(-self.time_warp_w + 1, self.time_warp_w)
|
||||
upper, lower = distorted[:w0, :], distorted[w0:, :]
|
||||
upper = cv2.resize(
|
||||
upper, dsize=(num_freqs, w0 + w), interpolation=cv2.INTER_LINEAR
|
||||
)
|
||||
lower = cv2.resize(
|
||||
lower,
|
||||
dsize=(num_freqs, num_frames - w0 - w),
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
)
|
||||
distorted = np.concatenate((upper, lower), axis=0)
|
||||
|
||||
for _i in range(self.freq_mask_n):
|
||||
f = np.random.randint(0, self.freq_mask_f)
|
||||
f0 = np.random.randint(0, num_freqs - f)
|
||||
if f != 0:
|
||||
distorted[:, f0 : f0 + f] = mask_value
|
||||
|
||||
max_time_mask_t = min(
|
||||
self.time_mask_t, math.floor(num_frames * self.time_mask_p)
|
||||
)
|
||||
if max_time_mask_t < 1:
|
||||
return distorted
|
||||
|
||||
for _i in range(self.time_mask_n):
|
||||
t = np.random.randint(0, max_time_mask_t)
|
||||
t0 = np.random.randint(0, num_frames - t)
|
||||
if t != 0:
|
||||
distorted[t0 : t0 + t, :] = mask_value
|
||||
|
||||
return distorted
|
||||
@@ -0,0 +1,40 @@
|
||||
import numpy as np
|
||||
from fairseq.data.audio.feature_transforms import (
|
||||
AudioFeatureTransform,
|
||||
register_audio_feature_transform,
|
||||
)
|
||||
|
||||
|
||||
@register_audio_feature_transform("utterance_cmvn")
|
||||
class UtteranceCMVN(AudioFeatureTransform):
|
||||
"""Utterance-level CMVN (cepstral mean and variance normalization)"""
|
||||
|
||||
@classmethod
|
||||
def from_config_dict(cls, config=None):
|
||||
_config = {} if config is None else config
|
||||
return UtteranceCMVN(
|
||||
_config.get("norm_means", True),
|
||||
_config.get("norm_vars", True),
|
||||
)
|
||||
|
||||
def __init__(self, norm_means=True, norm_vars=True):
|
||||
self.norm_means, self.norm_vars = norm_means, norm_vars
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
self.__class__.__name__
|
||||
+ f"(norm_means={self.norm_means}, norm_vars={self.norm_vars})"
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
mean = x.mean(axis=0)
|
||||
square_sums = (x ** 2).sum(axis=0)
|
||||
|
||||
if self.norm_means:
|
||||
x = np.subtract(x, mean)
|
||||
if self.norm_vars:
|
||||
var = square_sums / x.shape[0] - mean ** 2
|
||||
std = np.sqrt(np.maximum(var, 1e-10))
|
||||
x = np.divide(x, std)
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,176 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .. import FairseqDataset
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RawAudioDataset(FairseqDataset):
|
||||
def __init__(
|
||||
self,
|
||||
sample_rate,
|
||||
max_sample_size=None,
|
||||
min_sample_size=0,
|
||||
shuffle=True,
|
||||
pad=False,
|
||||
normalize=False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.sample_rate = sample_rate
|
||||
self.sizes = []
|
||||
self.max_sample_size = (
|
||||
max_sample_size if max_sample_size is not None else sys.maxsize
|
||||
)
|
||||
self.min_sample_size = min_sample_size
|
||||
self.pad = pad
|
||||
self.shuffle = shuffle
|
||||
self.normalize = normalize
|
||||
|
||||
def __getitem__(self, index):
|
||||
raise NotImplementedError()
|
||||
|
||||
def __len__(self):
|
||||
return len(self.sizes)
|
||||
|
||||
def postprocess(self, feats, curr_sample_rate):
|
||||
if feats.dim() == 2:
|
||||
feats = feats.mean(-1)
|
||||
|
||||
if curr_sample_rate != self.sample_rate:
|
||||
raise Exception(f"sample rate: {curr_sample_rate}, need {self.sample_rate}")
|
||||
|
||||
assert feats.dim() == 1, feats.dim()
|
||||
|
||||
if self.normalize:
|
||||
with torch.no_grad():
|
||||
feats = F.layer_norm(feats, feats.shape)
|
||||
return feats
|
||||
|
||||
def crop_to_max_size(self, wav, target_size):
|
||||
size = len(wav)
|
||||
diff = size - target_size
|
||||
if diff <= 0:
|
||||
return wav
|
||||
|
||||
start = np.random.randint(0, diff + 1)
|
||||
end = size - diff + start
|
||||
return wav[start:end]
|
||||
|
||||
def collater(self, samples):
|
||||
samples = [s for s in samples if s["source"] is not None]
|
||||
if len(samples) == 0:
|
||||
return {}
|
||||
|
||||
sources = [s["source"] for s in samples]
|
||||
sizes = [len(s) for s in sources]
|
||||
|
||||
if self.pad:
|
||||
target_size = min(max(sizes), self.max_sample_size)
|
||||
else:
|
||||
target_size = min(min(sizes), self.max_sample_size)
|
||||
|
||||
collated_sources = sources[0].new_zeros(len(sources), target_size)
|
||||
padding_mask = (
|
||||
torch.BoolTensor(collated_sources.shape).fill_(False) if self.pad else None
|
||||
)
|
||||
for i, (source, size) in enumerate(zip(sources, sizes)):
|
||||
diff = size - target_size
|
||||
if diff == 0:
|
||||
collated_sources[i] = source
|
||||
elif diff < 0:
|
||||
assert self.pad
|
||||
collated_sources[i] = torch.cat(
|
||||
[source, source.new_full((-diff,), 0.0)]
|
||||
)
|
||||
padding_mask[i, diff:] = True
|
||||
else:
|
||||
collated_sources[i] = self.crop_to_max_size(source, target_size)
|
||||
|
||||
input = {"source": collated_sources}
|
||||
if self.pad:
|
||||
input["padding_mask"] = padding_mask
|
||||
return {"id": torch.LongTensor([s["id"] for s in samples]), "net_input": input}
|
||||
|
||||
def num_tokens(self, index):
|
||||
return self.size(index)
|
||||
|
||||
def size(self, index):
|
||||
"""Return an example's size as a float or tuple. This value is used when
|
||||
filtering a dataset with ``--max-positions``."""
|
||||
if self.pad:
|
||||
return self.sizes[index]
|
||||
return min(self.sizes[index], self.max_sample_size)
|
||||
|
||||
def ordered_indices(self):
|
||||
"""Return an ordered list of indices. Batches will be constructed based
|
||||
on this order."""
|
||||
|
||||
if self.shuffle:
|
||||
order = [np.random.permutation(len(self))]
|
||||
else:
|
||||
order = [np.arange(len(self))]
|
||||
|
||||
order.append(self.sizes)
|
||||
return np.lexsort(order)[::-1]
|
||||
|
||||
|
||||
class FileAudioDataset(RawAudioDataset):
|
||||
def __init__(
|
||||
self,
|
||||
manifest_path,
|
||||
sample_rate,
|
||||
max_sample_size=None,
|
||||
min_sample_size=0,
|
||||
shuffle=True,
|
||||
pad=False,
|
||||
normalize=False,
|
||||
):
|
||||
super().__init__(
|
||||
sample_rate=sample_rate,
|
||||
max_sample_size=max_sample_size,
|
||||
min_sample_size=min_sample_size,
|
||||
shuffle=shuffle,
|
||||
pad=pad,
|
||||
normalize=normalize,
|
||||
)
|
||||
|
||||
self.fnames = []
|
||||
self.line_inds = set()
|
||||
|
||||
skipped = 0
|
||||
with open(manifest_path, "r") as f:
|
||||
self.root_dir = f.readline().strip()
|
||||
for i, line in enumerate(f):
|
||||
items = line.strip().split("\t")
|
||||
assert len(items) == 2, line
|
||||
sz = int(items[1])
|
||||
if min_sample_size is not None and sz < min_sample_size:
|
||||
skipped += 1
|
||||
continue
|
||||
self.fnames.append(items[0])
|
||||
self.line_inds.add(i)
|
||||
self.sizes.append(sz)
|
||||
logger.info(f"loaded {len(self.fnames)}, skipped {skipped} samples")
|
||||
|
||||
def __getitem__(self, index):
|
||||
import soundfile as sf
|
||||
|
||||
fname = os.path.join(self.root_dir, self.fnames[index])
|
||||
wav, curr_sample_rate = sf.read(fname)
|
||||
feats = torch.from_numpy(wav).float()
|
||||
feats = self.postprocess(feats, curr_sample_rate)
|
||||
return {"id": index, "source": feats}
|
||||
@@ -0,0 +1,528 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import csv
|
||||
import io
|
||||
import logging
|
||||
import os.path as op
|
||||
import re
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from fairseq.data import (
|
||||
ConcatDataset,
|
||||
Dictionary,
|
||||
FairseqDataset,
|
||||
ResamplingDataset,
|
||||
data_utils as fairseq_data_utils,
|
||||
)
|
||||
from fairseq.data.audio.audio_utils import get_fbank, get_waveform
|
||||
from fairseq.data.audio.feature_transforms import CompositeAudioFeatureTransform
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class S2TDataConfig(object):
|
||||
"""Wrapper class for data config YAML"""
|
||||
|
||||
def __init__(self, yaml_path):
|
||||
try:
|
||||
import yaml
|
||||
except ImportError:
|
||||
print("Please install PyYAML to load YAML files for " "S2T data config")
|
||||
self.config = {}
|
||||
if op.isfile(yaml_path):
|
||||
try:
|
||||
with open(yaml_path) as f:
|
||||
self.config = yaml.load(f, Loader=yaml.FullLoader)
|
||||
except Exception as e:
|
||||
logger.info(f"Failed to load config from {yaml_path}: {e}")
|
||||
else:
|
||||
logger.info(f"Cannot find {yaml_path}")
|
||||
|
||||
@property
|
||||
def vocab_filename(self):
|
||||
"""fairseq vocabulary file under data root"""
|
||||
return self.config.get("vocab_filename", "dict.txt")
|
||||
|
||||
@property
|
||||
def shuffle(self) -> bool:
|
||||
"""Shuffle dataset samples before batching"""
|
||||
return self.config.get("shuffle", False)
|
||||
|
||||
@property
|
||||
def pre_tokenizer(self) -> Dict:
|
||||
"""Pre-tokenizer to apply before subword tokenization. Returning
|
||||
a dictionary with `tokenizer` providing the tokenizer name and
|
||||
the other items providing the tokenizer-specific arguments.
|
||||
Tokenizers are defined in `fairseq.data.encoders.*`"""
|
||||
return self.config.get("pre_tokenizer", {"tokenizer": None})
|
||||
|
||||
@property
|
||||
def bpe_tokenizer(self) -> Dict:
|
||||
"""Subword tokenizer to apply after pre-tokenization. Returning
|
||||
a dictionary with `bpe` providing the tokenizer name and
|
||||
the other items providing the tokenizer-specific arguments.
|
||||
Tokenizers are defined in `fairseq.data.encoders.*`"""
|
||||
return self.config.get("bpe_tokenizer", {"bpe": None})
|
||||
|
||||
@property
|
||||
def prepend_tgt_lang_tag(self) -> bool:
|
||||
"""Prepend target lang ID token as the target BOS (e.g. for to-many
|
||||
multilingual setting). During inference, this requires `--prefix-size 1`
|
||||
to force BOS to be lang ID token."""
|
||||
return self.config.get("prepend_tgt_lang_tag", False)
|
||||
|
||||
@property
|
||||
def input_feat_per_channel(self):
|
||||
"""The dimension of input features (per audio channel)"""
|
||||
return self.config.get("input_feat_per_channel", 80)
|
||||
|
||||
@property
|
||||
def input_channels(self):
|
||||
"""The number of channels in the input audio"""
|
||||
return self.config.get("input_channels", 1)
|
||||
|
||||
@property
|
||||
def sampling_alpha(self):
|
||||
"""Hyper-parameter alpha = 1/T for temperature-based resampling.
|
||||
(alpha = 1 for no resampling)"""
|
||||
return self.config.get("sampling_alpha", 1.0)
|
||||
|
||||
@property
|
||||
def use_audio_input(self):
|
||||
"""Needed by the dataset loader to see if the model requires
|
||||
raw audio as inputs."""
|
||||
return self.config.get("use_audio_input", False)
|
||||
|
||||
@property
|
||||
def audio_root(self):
|
||||
"""Audio paths in the manifest TSV can be relative and this provides
|
||||
the root path. Set this to empty string when using absolute paths."""
|
||||
return self.config.get("audio_root", "")
|
||||
|
||||
def get_feature_transforms(self, split, is_train):
|
||||
"""Split-specific feature transforms. Allowing train set wildcard `_train`,
|
||||
evaluation set wildcard `_eval` and general wildcard `*` for matching."""
|
||||
from copy import deepcopy
|
||||
|
||||
cfg = deepcopy(self.config)
|
||||
_cur = cfg.get("transforms", {})
|
||||
cur = _cur.get(split)
|
||||
cur = _cur.get("_train") if cur is None and is_train else cur
|
||||
cur = _cur.get("_eval") if cur is None and not is_train else cur
|
||||
cur = _cur.get("*") if cur is None else cur
|
||||
cfg["transforms"] = cur
|
||||
return cfg
|
||||
|
||||
|
||||
def is_npy_data(data: bytes) -> bool:
|
||||
return data[0] == 147 and data[1] == 78
|
||||
|
||||
|
||||
def is_flac_or_wav_data(data: bytes) -> bool:
|
||||
is_flac = data[0] == 102 and data[1] == 76
|
||||
is_wav = data[0] == 82 and data[1] == 73
|
||||
return is_flac or is_wav
|
||||
|
||||
|
||||
def read_from_uncompressed_zip(file_path, offset, file_size) -> bytes:
|
||||
with open(file_path, "rb") as f:
|
||||
f.seek(offset)
|
||||
data = f.read(file_size)
|
||||
return data
|
||||
|
||||
|
||||
def get_features_from_npy_or_audio(path):
|
||||
ext = op.splitext(op.basename(path))[1]
|
||||
if ext not in {".npy", ".flac", ".wav"}:
|
||||
raise ValueError(f'Unsupported file format for "{path}"')
|
||||
return np.load(path) if ext == ".npy" else get_fbank(path)
|
||||
|
||||
|
||||
def get_features_or_waveform_from_uncompressed_zip(
|
||||
path, byte_offset, byte_size, need_waveform=False
|
||||
):
|
||||
assert path.endswith(".zip")
|
||||
data = read_from_uncompressed_zip(path, byte_offset, byte_size)
|
||||
f = io.BytesIO(data)
|
||||
if is_npy_data(data):
|
||||
features_or_waveform = np.load(f)
|
||||
elif is_flac_or_wav_data(data):
|
||||
features_or_waveform = get_waveform(f)[0] if need_waveform else get_fbank(f)
|
||||
else:
|
||||
raise ValueError(f'Unknown file format for "{path}"')
|
||||
return features_or_waveform
|
||||
|
||||
|
||||
def get_features_or_waveform(path: str, need_waveform=False):
|
||||
"""Get speech features from .npy file or waveform from .wav/.flac file.
|
||||
The file may be inside an uncompressed ZIP file and is accessed via byte
|
||||
offset and length.
|
||||
|
||||
Args:
|
||||
path (str): File path in the format of "<.npy/.wav/.flac path>" or
|
||||
"<zip path>:<byte offset>:<byte length>".
|
||||
need_waveform (bool): return waveform instead of features.
|
||||
|
||||
Returns:
|
||||
features_or_waveform (numpy.ndarray): speech features or waveform.
|
||||
"""
|
||||
_path, *extra = path.split(":")
|
||||
if not op.exists(_path):
|
||||
raise FileNotFoundError(f"File not found: {_path}")
|
||||
|
||||
if len(extra) == 0:
|
||||
if need_waveform:
|
||||
return get_waveform(_path)
|
||||
return get_features_from_npy_or_audio(_path)
|
||||
elif len(extra) == 2:
|
||||
extra = [int(i) for i in extra]
|
||||
features_or_waveform = get_features_or_waveform_from_uncompressed_zip(
|
||||
_path, extra[0], extra[1], need_waveform=need_waveform
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid path: {path}")
|
||||
|
||||
return features_or_waveform
|
||||
|
||||
|
||||
def _collate_frames(
|
||||
frames: List[torch.Tensor], is_audio_input: bool = False
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Convert a list of 2D frames into a padded 3D tensor
|
||||
Args:
|
||||
frames (list): list of 2D frames of size L[i]*f_dim. Where L[i] is
|
||||
length of i-th frame and f_dim is static dimension of features
|
||||
Returns:
|
||||
3D tensor of size len(frames)*len_max*f_dim where len_max is max of L[i]
|
||||
"""
|
||||
max_len = max(frame.size(0) for frame in frames)
|
||||
if is_audio_input:
|
||||
out = frames[0].new_zeros((len(frames), max_len))
|
||||
else:
|
||||
out = frames[0].new_zeros((len(frames), max_len, frames[0].size(1)))
|
||||
for i, v in enumerate(frames):
|
||||
out[i, : v.size(0)] = v
|
||||
return out
|
||||
|
||||
|
||||
class SpeechToTextDataset(FairseqDataset):
|
||||
LANG_TAG_TEMPLATE = "<lang:{}>"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
split: str,
|
||||
is_train_split: bool,
|
||||
data_cfg: S2TDataConfig,
|
||||
audio_paths: List[str],
|
||||
n_frames: List[int],
|
||||
src_texts: Optional[List[str]] = None,
|
||||
tgt_texts: Optional[List[str]] = None,
|
||||
speakers: Optional[List[str]] = None,
|
||||
src_langs: Optional[List[str]] = None,
|
||||
tgt_langs: Optional[List[str]] = None,
|
||||
ids: Optional[List[str]] = None,
|
||||
tgt_dict: Optional[Dictionary] = None,
|
||||
pre_tokenizer=None,
|
||||
bpe_tokenizer=None,
|
||||
):
|
||||
self.split, self.is_train_split = split, is_train_split
|
||||
self.data_cfg = data_cfg
|
||||
self.audio_paths, self.n_frames = audio_paths, n_frames
|
||||
self.n_samples = len(audio_paths)
|
||||
assert len(n_frames) == self.n_samples > 0
|
||||
assert src_texts is None or len(src_texts) == self.n_samples
|
||||
assert tgt_texts is None or len(tgt_texts) == self.n_samples
|
||||
assert speakers is None or len(speakers) == self.n_samples
|
||||
assert src_langs is None or len(src_langs) == self.n_samples
|
||||
assert tgt_langs is None or len(tgt_langs) == self.n_samples
|
||||
assert ids is None or len(ids) == self.n_samples
|
||||
assert (tgt_dict is None and tgt_texts is None) or (
|
||||
tgt_dict is not None and tgt_texts is not None
|
||||
)
|
||||
self.src_texts, self.tgt_texts = src_texts, tgt_texts
|
||||
self.src_langs, self.tgt_langs = src_langs, tgt_langs
|
||||
self.tgt_dict = tgt_dict
|
||||
self.check_tgt_lang_tag()
|
||||
self.ids = ids
|
||||
self.shuffle = data_cfg.shuffle if is_train_split else False
|
||||
|
||||
self.feature_transforms = CompositeAudioFeatureTransform.from_config_dict(
|
||||
self.data_cfg.get_feature_transforms(split, is_train_split)
|
||||
)
|
||||
|
||||
self.pre_tokenizer = pre_tokenizer
|
||||
self.bpe_tokenizer = bpe_tokenizer
|
||||
|
||||
logger.info(self.__repr__())
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
self.__class__.__name__
|
||||
+ f'(split="{self.split}", n_samples={self.n_samples}, '
|
||||
f"prepend_tgt_lang_tag={self.data_cfg.prepend_tgt_lang_tag}, "
|
||||
f"shuffle={self.shuffle}, transforms={self.feature_transforms})"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def is_lang_tag(cls, token):
|
||||
pattern = cls.LANG_TAG_TEMPLATE.replace("{}", "(.*)")
|
||||
return re.match(pattern, token)
|
||||
|
||||
def check_tgt_lang_tag(self):
|
||||
if self.data_cfg.prepend_tgt_lang_tag:
|
||||
assert self.tgt_langs is not None and self.tgt_dict is not None
|
||||
tgt_lang_tags = [
|
||||
self.LANG_TAG_TEMPLATE.format(t) for t in set(self.tgt_langs)
|
||||
]
|
||||
assert all(t in self.tgt_dict for t in tgt_lang_tags)
|
||||
|
||||
def tokenize_text(self, text: str):
|
||||
if self.pre_tokenizer is not None:
|
||||
text = self.pre_tokenizer.encode(text)
|
||||
if self.bpe_tokenizer is not None:
|
||||
text = self.bpe_tokenizer.encode(text)
|
||||
return text
|
||||
|
||||
def __getitem__(
|
||||
self, index: int
|
||||
) -> Tuple[int, torch.Tensor, Optional[torch.Tensor]]:
|
||||
source = get_features_or_waveform(
|
||||
self.audio_paths[index], need_waveform=self.data_cfg.use_audio_input
|
||||
)
|
||||
if self.feature_transforms is not None:
|
||||
assert not self.data_cfg.use_audio_input
|
||||
source = self.feature_transforms(source)
|
||||
source = torch.from_numpy(source).float()
|
||||
|
||||
target = None
|
||||
if self.tgt_texts is not None:
|
||||
tokenized = self.tokenize_text(self.tgt_texts[index])
|
||||
target = self.tgt_dict.encode_line(
|
||||
tokenized, add_if_not_exist=False, append_eos=True
|
||||
).long()
|
||||
if self.data_cfg.prepend_tgt_lang_tag:
|
||||
lang_tag = self.LANG_TAG_TEMPLATE.format(self.tgt_langs[index])
|
||||
lang_tag_idx = self.tgt_dict.index(lang_tag)
|
||||
target = torch.cat((torch.LongTensor([lang_tag_idx]), target), 0)
|
||||
return index, source, target
|
||||
|
||||
def __len__(self):
|
||||
return self.n_samples
|
||||
|
||||
def collater(self, samples: List[Tuple[int, torch.Tensor, torch.Tensor]]) -> Dict:
|
||||
if len(samples) == 0:
|
||||
return {}
|
||||
indices = torch.tensor([i for i, _, _ in samples], dtype=torch.long)
|
||||
frames = _collate_frames(
|
||||
[s for _, s, _ in samples], self.data_cfg.use_audio_input
|
||||
)
|
||||
# sort samples by descending number of frames
|
||||
n_frames = torch.tensor([s.size(0) for _, s, _ in samples], dtype=torch.long)
|
||||
n_frames, order = n_frames.sort(descending=True)
|
||||
indices = indices.index_select(0, order)
|
||||
frames = frames.index_select(0, order)
|
||||
|
||||
target, target_lengths = None, None
|
||||
prev_output_tokens = None
|
||||
ntokens = None
|
||||
if self.tgt_texts is not None:
|
||||
target = fairseq_data_utils.collate_tokens(
|
||||
[t for _, _, t in samples],
|
||||
self.tgt_dict.pad(),
|
||||
self.tgt_dict.eos(),
|
||||
left_pad=False,
|
||||
move_eos_to_beginning=False,
|
||||
)
|
||||
target = target.index_select(0, order)
|
||||
target_lengths = torch.tensor(
|
||||
[t.size(0) for _, _, t in samples], dtype=torch.long
|
||||
).index_select(0, order)
|
||||
prev_output_tokens = fairseq_data_utils.collate_tokens(
|
||||
[t for _, _, t in samples],
|
||||
self.tgt_dict.pad(),
|
||||
self.tgt_dict.eos(),
|
||||
left_pad=False,
|
||||
move_eos_to_beginning=True,
|
||||
)
|
||||
prev_output_tokens = prev_output_tokens.index_select(0, order)
|
||||
ntokens = sum(t.size(0) for _, _, t in samples)
|
||||
|
||||
out = {
|
||||
"id": indices,
|
||||
"net_input": {
|
||||
"src_tokens": frames,
|
||||
"src_lengths": n_frames,
|
||||
"prev_output_tokens": prev_output_tokens,
|
||||
},
|
||||
"target": target,
|
||||
"target_lengths": target_lengths,
|
||||
"ntokens": ntokens,
|
||||
"nsentences": len(samples),
|
||||
}
|
||||
return out
|
||||
|
||||
def num_tokens(self, index):
|
||||
return self.n_frames[index]
|
||||
|
||||
def size(self, index):
|
||||
t_len = 0
|
||||
if self.tgt_texts is not None:
|
||||
tokenized = self.tokenize_text(self.tgt_texts[index])
|
||||
t_len = len(tokenized.split(" "))
|
||||
return self.n_frames[index], t_len
|
||||
|
||||
@property
|
||||
def sizes(self):
|
||||
return np.array(self.n_frames)
|
||||
|
||||
@property
|
||||
def can_reuse_epoch_itr_across_epochs(self):
|
||||
return True
|
||||
|
||||
def ordered_indices(self):
|
||||
if self.shuffle:
|
||||
order = [np.random.permutation(len(self))]
|
||||
else:
|
||||
order = [np.arange(len(self))]
|
||||
# first by descending order of # of frames then by original/random order
|
||||
order.append([-n for n in self.n_frames])
|
||||
return np.lexsort(order)
|
||||
|
||||
def prefetch(self, indices):
|
||||
raise False
|
||||
|
||||
|
||||
class SpeechToTextDatasetCreator(object):
|
||||
# mandatory columns
|
||||
KEY_ID, KEY_AUDIO, KEY_N_FRAMES = "id", "audio", "n_frames"
|
||||
KEY_TGT_TEXT = "tgt_text"
|
||||
# optional columns
|
||||
KEY_SPEAKER, KEY_SRC_TEXT = "speaker", "src_text"
|
||||
KEY_SRC_LANG, KEY_TGT_LANG = "src_lang", "tgt_lang"
|
||||
# default values
|
||||
DEFAULT_SPEAKER = DEFAULT_SRC_TEXT = DEFAULT_LANG = ""
|
||||
|
||||
@classmethod
|
||||
def _from_list(
|
||||
cls,
|
||||
split_name: str,
|
||||
is_train_split,
|
||||
samples: List[List[Dict]],
|
||||
data_cfg: S2TDataConfig,
|
||||
tgt_dict,
|
||||
pre_tokenizer,
|
||||
bpe_tokenizer,
|
||||
) -> SpeechToTextDataset:
|
||||
audio_paths, n_frames, src_texts, tgt_texts, ids = [], [], [], [], []
|
||||
speakers, src_langs, tgt_langs = [], [], []
|
||||
for s in samples:
|
||||
ids.extend([ss[cls.KEY_ID] for ss in s])
|
||||
audio_paths.extend(
|
||||
[op.join(data_cfg.audio_root, ss[cls.KEY_AUDIO]) for ss in s]
|
||||
)
|
||||
n_frames.extend([int(ss[cls.KEY_N_FRAMES]) for ss in s])
|
||||
tgt_texts.extend([ss[cls.KEY_TGT_TEXT] for ss in s])
|
||||
src_texts.extend(
|
||||
[ss.get(cls.KEY_SRC_TEXT, cls.DEFAULT_SRC_TEXT) for ss in s]
|
||||
)
|
||||
speakers.extend([ss.get(cls.KEY_SPEAKER, cls.DEFAULT_SPEAKER) for ss in s])
|
||||
src_langs.extend([ss.get(cls.KEY_SRC_LANG, cls.DEFAULT_LANG) for ss in s])
|
||||
tgt_langs.extend([ss.get(cls.KEY_TGT_LANG, cls.DEFAULT_LANG) for ss in s])
|
||||
return SpeechToTextDataset(
|
||||
split_name,
|
||||
is_train_split,
|
||||
data_cfg,
|
||||
audio_paths,
|
||||
n_frames,
|
||||
src_texts,
|
||||
tgt_texts,
|
||||
speakers,
|
||||
src_langs,
|
||||
tgt_langs,
|
||||
ids,
|
||||
tgt_dict,
|
||||
pre_tokenizer,
|
||||
bpe_tokenizer,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _get_size_ratios(cls, ids: List[str], sizes: List[int], alpha: float = 1.0):
|
||||
"""Size ratios for temperature-based sampling
|
||||
(https://arxiv.org/abs/1907.05019)"""
|
||||
_sizes = np.array(sizes)
|
||||
prob = _sizes / _sizes.sum()
|
||||
smoothed_prob = prob ** alpha
|
||||
smoothed_prob = smoothed_prob / smoothed_prob.sum()
|
||||
size_ratio = (smoothed_prob * _sizes.sum()) / _sizes
|
||||
|
||||
o_str = str({_i: f"{prob[i]:.3f}" for i, _i in enumerate(ids)})
|
||||
logger.info(f"original sampling probability: {o_str}")
|
||||
p_str = str({_i: f"{smoothed_prob[i]:.3f}" for i, _i in enumerate(ids)})
|
||||
logger.info(f"balanced sampling probability: {p_str}")
|
||||
sr_str = str({_id: f"{size_ratio[i]:.3f}" for i, _id in enumerate(ids)})
|
||||
logger.info(f"balanced sampling size ratio: {sr_str}")
|
||||
return size_ratio.tolist()
|
||||
|
||||
@classmethod
|
||||
def from_tsv(
|
||||
cls,
|
||||
root: str,
|
||||
data_cfg: S2TDataConfig,
|
||||
splits: str,
|
||||
tgt_dict,
|
||||
pre_tokenizer,
|
||||
bpe_tokenizer,
|
||||
is_train_split: bool,
|
||||
epoch: int,
|
||||
seed: int,
|
||||
) -> SpeechToTextDataset:
|
||||
samples = []
|
||||
_splits = splits.split(",")
|
||||
for split in _splits:
|
||||
tsv_path = op.join(root, f"{split}.tsv")
|
||||
if not op.isfile(tsv_path):
|
||||
raise FileNotFoundError(f"Dataset not found: {tsv_path}")
|
||||
with open(tsv_path) as f:
|
||||
reader = csv.DictReader(
|
||||
f,
|
||||
delimiter="\t",
|
||||
quotechar=None,
|
||||
doublequote=False,
|
||||
lineterminator="\n",
|
||||
quoting=csv.QUOTE_NONE,
|
||||
)
|
||||
samples.append([dict(e) for e in reader])
|
||||
assert len(samples) > 0
|
||||
|
||||
datasets = [
|
||||
cls._from_list(
|
||||
name,
|
||||
is_train_split,
|
||||
[s],
|
||||
data_cfg,
|
||||
tgt_dict,
|
||||
pre_tokenizer,
|
||||
bpe_tokenizer,
|
||||
)
|
||||
for name, s in zip(_splits, samples)
|
||||
]
|
||||
|
||||
if is_train_split and len(_splits) > 1 and data_cfg.sampling_alpha != 1.0:
|
||||
# temperature-based sampling
|
||||
size_ratios = cls._get_size_ratios(
|
||||
_splits, [len(s) for s in samples], alpha=data_cfg.sampling_alpha
|
||||
)
|
||||
datasets = [
|
||||
ResamplingDataset(
|
||||
d, size_ratio=r, seed=seed, epoch=epoch, replace=(r >= 1.0)
|
||||
)
|
||||
for d, r in zip(datasets, size_ratios)
|
||||
]
|
||||
return ConcatDataset(datasets)
|
||||
Reference in New Issue
Block a user