diff --git a/docs/README.md b/docs/README.md index c92987f..7da1a19 100644 --- a/docs/README.md +++ b/docs/README.md @@ -16,3 +16,4 @@ pip install -r requirements.txt * Data & Data Pipeline Design * [Data Format & Prepare](./data_format.md) * [Data pack(sequence pack)](./data_pack.md) + * [Data Extractor](./data_extractor.md) diff --git a/docs/data_extractor.md b/docs/data_extractor.md new file mode 100644 index 0000000..8f5b3d4 --- /dev/null +++ b/docs/data_extractor.md @@ -0,0 +1,86 @@ +## Data Extactor + +We designed Model Extractor for each model to extract the inputs required by that model. +Each Model Extractor inherits from the `Extractor` defined by WEST, as shown below: + +``` python +class Extractor(ABC): + + model_type = 'model' + + # Batch/Pack fileds for dataset + fields_batch_static = {} + fields_batch_dynamic = {} + fields_pack_offset = {} + + def __init__(self, tokenizer, model_config, inference=False): + self.tokenizer = tokenizer + self.model_config = model_config + self.inference = inference + + @abstractmethod + def extract(self, item): + pass +``` + +* `fields_batch_static`: fileds which has static shape, such as text lengths, speech feature lengths. +* `fields_batch_dynamic`: fileds which has dynamic shape, which requires padding when batching, such as speech feature. +* `fields_pack_offset`: fileds we want to keep offset record when packing sequence, such as the speech token start index. + +The `extract` function returns a dict object, and here `fileds` are keys in the returned dict. + + +Here is an example of `TouchTTS` Extactor. + + +``` python + +class ExtractorTouchTTS(Extractor): + model_type = 'touch_tts' + fields_batch_static = {'audio_offsets', 'text_lengths'} + fields_batch_dynamic = {'audio_features', 'input_ids', 'labels'} + fields_pack_offset = {'audio_offsets'} + + def extract(self, item): + import s3tokenizer + IGNORE_TOKEN_ID = LabelSmoother.ignore_index + waveform, sample_rate = torchaudio.load(item['wav']) + audio = torchaudio.transforms.Resample(sample_rate, 16000)(waveform) + audio = audio[0] # get the first channel + mel = s3tokenizer.log_mel_spectrogram(audio) + mel = mel.transpose(0, 1) + # There is 100 frames mel per second, and 25 tokens per second + num_audio_token = math.ceil(mel.size(0) * 25 / 100.0 - 1e-9) + if not self.inference: + content = item['txt'] + '<|audio_bos|>' + token_lengths = 0 + else: + content = item['txt'] + item['syn_txt'] + '<|audio_bos|>' + token_lengths = len(self.tokenizer.encode(item['syn_txt'])) + ids_text = [self.tokenizer.bos_token_id + ] + self.tokenizer.encode(content) + tgt_text = [IGNORE_TOKEN_ID] * len(ids_text) + ids_audio = [0] * num_audio_token + if not self.inference: + ids = ids_text + ids_audio + [self.tokenizer.eos_token_id] + tgt = tgt_text + ids_audio + [self.tokenizer.eos_token_id] + else: + ids = ids_text + ids_audio + tgt = tgt_text + ids_audio + input_ids = torch.tensor(ids, dtype=torch.long) + tgt_ids = torch.tensor(tgt, dtype=torch.long) + return { + 'input_ids': input_ids, + 'labels': tgt_ids, + 'audio_features': mel, + 'audio_offsets': len(ids_text), + 'text_lengths': token_lengths + } +``` + + +The `ExtractorTouchTTS` returns a dict and in which: + +* `input_ids`, `labels`, `audio_features` are in dyanmic shape. +* `audio_offsets`, `text_lengths` are in static shape. +* `audio_offsets` is the speech token start offset.