pheme / data /semantic_dataset.py
taras-sereda's picture
minimal set of files to run inference; pheme-small checkpoint
96ee597
raw
history blame
7.09 kB
"""Semantic tokens loading logic.
Copyright PolyAI Limited.
"""
import json
import logging
import random
import re
from logging import getLogger
from pathlib import Path
from typing import List, Pattern, Union
import numpy as np
import torch
from phonemizer.backend import EspeakBackend
from phonemizer.backend.espeak.language_switch import LanguageSwitch
from phonemizer.backend.espeak.words_mismatch import WordMismatch
from phonemizer.punctuation import Punctuation
from phonemizer.separator import Separator
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm
from data.collation import get_text_semantic_token_collater
class TextTokenizer:
"""Phonemize Text."""
def __init__(
self,
language="en-us",
backend="espeak",
separator=Separator(word="_", syllable="-", phone="|"),
preserve_punctuation=True,
punctuation_marks: Union[str, Pattern] = Punctuation.default_marks(),
with_stress: bool = False,
tie: Union[bool, str] = False,
language_switch: LanguageSwitch = "keep-flags",
words_mismatch: WordMismatch = "ignore",
) -> None:
logger = getLogger("phonemizer")
logger.setLevel(logging.ERROR)
if backend == "espeak":
phonemizer = EspeakBackend(
language,
punctuation_marks=punctuation_marks,
preserve_punctuation=preserve_punctuation,
with_stress=with_stress,
tie=tie,
language_switch=language_switch,
words_mismatch=words_mismatch,
logger=logger,
)
else:
raise NotImplementedError(f"{backend}")
self.backend = phonemizer
self.separator = separator
def to_list(self, phonemized: str) -> List[str]:
fields = []
for word in phonemized.split(self.separator.word):
# "ɐ m|iː|n?" ɹ|ɪ|z|ɜː|v; h|ɪ|z.
pp = re.findall(r"\w+|[^\w\s]", word, re.UNICODE)
fields.extend(
[p for p in pp if p != self.separator.phone] + [self.separator.word]
)
assert len("".join(fields[:-1])) == len(phonemized) - phonemized.count(
self.separator.phone
)
return fields[:-1]
def __call__(self, text, strip=True) -> List[List[str]]:
if isinstance(text, str):
text = [text]
phonemized = self.backend.phonemize(
text, separator=self.separator, strip=strip, njobs=1
)
return [self.to_list(p) for p in phonemized]
class Collator:
def collate(self, batch):
input_ids = [item["input_ids"] for item in batch]
output_sequences = [item["labels"] for item in batch]
# Pad sequences to the maximum length in the batch
input_ids = torch.nn.utils.rnn.pad_sequence(
input_ids, batch_first=True, padding_value=0
)
output_sequences = torch.nn.utils.rnn.pad_sequence(
output_sequences, batch_first=True, padding_value=-100
)
# 1 - token is unmasked, 0 - token is masked.
attention_mask = input_ids != 0
return {
"input_ids": input_ids,
"attention_mask": attention_mask,
"labels": output_sequences,
}
class ConcatenateSemanticDataset(Dataset):
def __init__(
self, manifest_path: str, symbol_table_path: str,
n_samples: int = 0, max_duration=15):
self.data = []
self.phonemizer = TextTokenizer()
self.text_collater = get_text_semantic_token_collater(
symbol_table_path)
self.manifest_path = manifest_path
self.n_samples = n_samples
self.max_duration = max_duration
if manifest_path is not None:
self._build()
def __len__(self):
if self.n_samples:
return min(self.n_samples, len(self.data))
return len(self.data)
def remove_unknown_symbols(self, text: List[str]):
res = []
for sym in text:
if sym not in self.text_collater.token2idx:
# print(f'{sym} is unk')
continue
res.append(sym)
return res
def __getitem__(self, idx):
item = self.data[idx]
input_ids = item["phoneme"].split("|")
input_ids = self.remove_unknown_symbols(input_ids)
input_ids_2 = None
if item.get("phoneme_2"):
input_ids_2 = item["phoneme_2"].split("|")
input_ids_2 = [self.remove_unknown_symbols(input_ids_2)]
input_ids = self.text_collater(
[input_ids], input_ids_2).to(dtype=torch.long)
input_ids = input_ids.to(dtype=torch.long)
labels = np.load(item["semantic_path"])
labels = [str(lbl) for lbl in labels]
labels_2 = None
if item.get("semantic_path_2"):
labels_2 = np.load(item["semantic_path_2"])
labels_2 = [[str(lbl) for lbl in labels_2]]
labels = self.text_collater([labels], labels_2).to(dtype=torch.long)
return {"input_ids": input_ids.squeeze(0), "labels": labels.squeeze(0)}
# TODO - remove this to not load to the memory
def _build(self):
for manifest_path in self.manifest_path:
dataset_path = Path(manifest_path).parent
with open(manifest_path, "r") as manifest_file:
manifest_data = json.load(manifest_file)
for key, value in tqdm(manifest_data.items()):
if float(value["duration"]) > self.max_duration:
continue
text = value["text"]
phoneme = value["phoneme"]
npy_path = f"{dataset_path}/audios-speech-tokenizer/semantic/{key.split('.wav')[0]}.npy" # noqa
datapoint = {
"text": text,
"semantic_path": npy_path,
"phoneme": phoneme
}
self.data.append(datapoint)
print(f"Total length of the dataset {manifest_path}: {len(self.data)}")
random.shuffle(self.data)
if __name__ == "__main__":
# Create an instance of the dataset
manifest_path = "datasets/ljspeech-training-data/dev.json"
text_tokens_file = "ckpt/unique_text_tokens.k2symbols"
seq2seq_dataset = ConcatenateSemanticDataset(
[manifest_path, manifest_path], text_tokens_file)
# seq2seq_dataset.phonemize_and_rewrite_manifest()
batch_size = 1 # Adjust to your desired batch size
dataloader = DataLoader(
seq2seq_dataset,
batch_size=batch_size,
shuffle=True,
collate_fn=Collator().collate,
)
for batch in dataloader:
print(batch["input_ids"])
print(batch["labels"])
print(batch["input_ids"][0].unique().max())
print(batch["input_ids"][0].unique().min())
print(batch["input_ids"].shape)
print(batch["labels"].shape)
break # Stop after the first batch if needed