502 lines
16 KiB
Python
502 lines
16 KiB
Python
|
from dataclasses import replace
|
|||
|
|
from typing import Dict
|
|||
|
|
|
|||
|
|
from TTS.tts.configs.shared_configs import CharactersConfig
|
|||
|
|
|
|||
|
|
|
|||
|
|
def parse_symbols():
|
|||
|
|
return {
|
|||
|
|
"pad": _pad,
|
|||
|
|
"eos": _eos,
|
|||
|
|
"bos": _bos,
|
|||
|
|
"characters": _characters,
|
|||
|
|
"punctuations": _punctuations,
|
|||
|
|
"phonemes": _phonemes,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
# DEFAULT SET OF GRAPHEMES
|
|||
|
|
_pad = "<PAD>"
|
|||
|
|
_eos = "<EOS>"
|
|||
|
|
_bos = "<BOS>"
|
|||
|
|
_blank = "<BLNK>" # TODO: check if we need this alongside with PAD
|
|||
|
|
_characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
|||
|
|
_punctuations = "!'(),-.:;? "
|
|||
|
|
|
|||
|
|
|
|||
|
|
# DEFAULT SET OF IPA PHONEMES
|
|||
|
|
# Phonemes definition (All IPA characters)
|
|||
|
|
_vowels = "iyɨʉɯuɪʏʊeøɘəɵɤoɛœɜɞʌɔæɐaɶɑɒᵻ"
|
|||
|
|
_non_pulmonic_consonants = "ʘɓǀɗǃʄǂɠǁʛ"
|
|||
|
|
_pulmonic_consonants = "pbtdʈɖcɟkɡqɢʔɴŋɲɳnɱmʙrʀⱱɾɽɸβfvθðszʃʒʂʐçʝxɣχʁħʕhɦɬɮʋɹɻjɰlɭʎʟ"
|
|||
|
|
_suprasegmentals = "ˈˌːˑ"
|
|||
|
|
_other_symbols = "ʍwɥʜʢʡɕʑɺɧʲ"
|
|||
|
|
_diacrilics = "ɚ˞ɫ"
|
|||
|
|
_phonemes = _vowels + _non_pulmonic_consonants + _pulmonic_consonants + _suprasegmentals + _other_symbols + _diacrilics
|
|||
|
|
|
|||
|
|
|
|||
|
|
class BaseVocabulary:
|
|||
|
|
"""Base Vocabulary class.
|
|||
|
|
|
|||
|
|
This class only needs a vocabulary dictionary without specifying the characters.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
vocab (Dict): A dictionary of characters and their corresponding indices.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, vocab: Dict, pad: str = None, blank: str = None, bos: str = None, eos: str = None):
|
|||
|
|
self.vocab = vocab
|
|||
|
|
self.pad = pad
|
|||
|
|
self.blank = blank
|
|||
|
|
self.bos = bos
|
|||
|
|
self.eos = eos
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def pad_id(self) -> int:
|
|||
|
|
"""Return the index of the padding character. If the padding character is not specified, return the length
|
|||
|
|
of the vocabulary."""
|
|||
|
|
return self.char_to_id(self.pad) if self.pad else len(self.vocab)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def blank_id(self) -> int:
|
|||
|
|
"""Return the index of the blank character. If the blank character is not specified, return the length of
|
|||
|
|
the vocabulary."""
|
|||
|
|
return self.char_to_id(self.blank) if self.blank else len(self.vocab)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def bos_id(self) -> int:
|
|||
|
|
"""Return the index of the bos character. If the bos character is not specified, return the length of the
|
|||
|
|
vocabulary."""
|
|||
|
|
return self.char_to_id(self.bos) if self.bos else len(self.vocab)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def eos_id(self) -> int:
|
|||
|
|
"""Return the index of the eos character. If the eos character is not specified, return the length of the
|
|||
|
|
vocabulary."""
|
|||
|
|
return self.char_to_id(self.eos) if self.eos else len(self.vocab)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def vocab(self):
|
|||
|
|
"""Return the vocabulary dictionary."""
|
|||
|
|
return self._vocab
|
|||
|
|
|
|||
|
|
@vocab.setter
|
|||
|
|
def vocab(self, vocab):
|
|||
|
|
"""Set the vocabulary dictionary and character mapping dictionaries."""
|
|||
|
|
self._vocab, self._char_to_id, self._id_to_char = None, None, None
|
|||
|
|
if vocab is not None:
|
|||
|
|
self._vocab = vocab
|
|||
|
|
self._char_to_id = {char: idx for idx, char in enumerate(self._vocab)}
|
|||
|
|
self._id_to_char = {
|
|||
|
|
idx: char for idx, char in enumerate(self._vocab) # pylint: disable=unnecessary-comprehension
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def init_from_config(config, **kwargs):
|
|||
|
|
"""Initialize from the given config."""
|
|||
|
|
if config.characters is not None and "vocab_dict" in config.characters and config.characters.vocab_dict:
|
|||
|
|
return (
|
|||
|
|
BaseVocabulary(
|
|||
|
|
config.characters.vocab_dict,
|
|||
|
|
config.characters.pad,
|
|||
|
|
config.characters.blank,
|
|||
|
|
config.characters.bos,
|
|||
|
|
config.characters.eos,
|
|||
|
|
),
|
|||
|
|
config,
|
|||
|
|
)
|
|||
|
|
return BaseVocabulary(**kwargs), config
|
|||
|
|
|
|||
|
|
def to_config(self) -> "CharactersConfig":
|
|||
|
|
return CharactersConfig(
|
|||
|
|
vocab_dict=self._vocab,
|
|||
|
|
pad=self.pad,
|
|||
|
|
eos=self.eos,
|
|||
|
|
bos=self.bos,
|
|||
|
|
blank=self.blank,
|
|||
|
|
is_unique=False,
|
|||
|
|
is_sorted=False,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def num_chars(self):
|
|||
|
|
"""Return number of tokens in the vocabulary."""
|
|||
|
|
return len(self._vocab)
|
|||
|
|
|
|||
|
|
def char_to_id(self, char: str) -> int:
|
|||
|
|
"""Map a character to an token ID."""
|
|||
|
|
try:
|
|||
|
|
return self._char_to_id[char]
|
|||
|
|
except KeyError as e:
|
|||
|
|
raise KeyError(f" [!] {repr(char)} is not in the vocabulary.") from e
|
|||
|
|
|
|||
|
|
def id_to_char(self, idx: int) -> str:
|
|||
|
|
"""Map an token ID to a character."""
|
|||
|
|
return self._id_to_char[idx]
|
|||
|
|
|
|||
|
|
|
|||
|
|
class BaseCharacters:
|
|||
|
|
"""🐸BaseCharacters class
|
|||
|
|
|
|||
|
|
Every new character class should inherit from this.
|
|||
|
|
|
|||
|
|
Characters are oredered as follows ```[PAD, EOS, BOS, BLANK, CHARACTERS, PUNCTUATIONS]```.
|
|||
|
|
|
|||
|
|
If you need a custom order, you need to define inherit from this class and override the ```_create_vocab``` method.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
characters (str):
|
|||
|
|
Main set of characters to be used in the vocabulary.
|
|||
|
|
|
|||
|
|
punctuations (str):
|
|||
|
|
Characters to be treated as punctuation.
|
|||
|
|
|
|||
|
|
pad (str):
|
|||
|
|
Special padding character that would be ignored by the model.
|
|||
|
|
|
|||
|
|
eos (str):
|
|||
|
|
End of the sentence character.
|
|||
|
|
|
|||
|
|
bos (str):
|
|||
|
|
Beginning of the sentence character.
|
|||
|
|
|
|||
|
|
blank (str):
|
|||
|
|
Optional character used between characters by some models for better prosody.
|
|||
|
|
|
|||
|
|
is_unique (bool):
|
|||
|
|
Remove duplicates from the provided characters. Defaults to True.
|
|||
|
|
el
|
|||
|
|
is_sorted (bool):
|
|||
|
|
Sort the characters in alphabetical order. Only applies to `self.characters`. Defaults to True.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
characters: str = None,
|
|||
|
|
punctuations: str = None,
|
|||
|
|
pad: str = None,
|
|||
|
|
eos: str = None,
|
|||
|
|
bos: str = None,
|
|||
|
|
blank: str = None,
|
|||
|
|
is_unique: bool = False,
|
|||
|
|
is_sorted: bool = True,
|
|||
|
|
) -> None:
|
|||
|
|
self._characters = characters
|
|||
|
|
self._punctuations = punctuations
|
|||
|
|
self._pad = pad
|
|||
|
|
self._eos = eos
|
|||
|
|
self._bos = bos
|
|||
|
|
self._blank = blank
|
|||
|
|
self.is_unique = is_unique
|
|||
|
|
self.is_sorted = is_sorted
|
|||
|
|
self._create_vocab()
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def pad_id(self) -> int:
|
|||
|
|
return self.char_to_id(self.pad) if self.pad else len(self.vocab)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def blank_id(self) -> int:
|
|||
|
|
return self.char_to_id(self.blank) if self.blank else len(self.vocab)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def eos_id(self) -> int:
|
|||
|
|
return self.char_to_id(self.eos) if self.eos else len(self.vocab)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def bos_id(self) -> int:
|
|||
|
|
return self.char_to_id(self.bos) if self.bos else len(self.vocab)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def characters(self):
|
|||
|
|
return self._characters
|
|||
|
|
|
|||
|
|
@characters.setter
|
|||
|
|
def characters(self, characters):
|
|||
|
|
self._characters = characters
|
|||
|
|
self._create_vocab()
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def punctuations(self):
|
|||
|
|
return self._punctuations
|
|||
|
|
|
|||
|
|
@punctuations.setter
|
|||
|
|
def punctuations(self, punctuations):
|
|||
|
|
self._punctuations = punctuations
|
|||
|
|
self._create_vocab()
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def pad(self):
|
|||
|
|
return self._pad
|
|||
|
|
|
|||
|
|
@pad.setter
|
|||
|
|
def pad(self, pad):
|
|||
|
|
self._pad = pad
|
|||
|
|
self._create_vocab()
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def eos(self):
|
|||
|
|
return self._eos
|
|||
|
|
|
|||
|
|
@eos.setter
|
|||
|
|
def eos(self, eos):
|
|||
|
|
self._eos = eos
|
|||
|
|
self._create_vocab()
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def bos(self):
|
|||
|
|
return self._bos
|
|||
|
|
|
|||
|
|
@bos.setter
|
|||
|
|
def bos(self, bos):
|
|||
|
|
self._bos = bos
|
|||
|
|
self._create_vocab()
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def blank(self):
|
|||
|
|
return self._blank
|
|||
|
|
|
|||
|
|
@blank.setter
|
|||
|
|
def blank(self, blank):
|
|||
|
|
self._blank = blank
|
|||
|
|
self._create_vocab()
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def vocab(self):
|
|||
|
|
return self._vocab
|
|||
|
|
|
|||
|
|
@vocab.setter
|
|||
|
|
def vocab(self, vocab):
|
|||
|
|
self._vocab = vocab
|
|||
|
|
self._char_to_id = {char: idx for idx, char in enumerate(self.vocab)}
|
|||
|
|
self._id_to_char = {
|
|||
|
|
idx: char for idx, char in enumerate(self.vocab) # pylint: disable=unnecessary-comprehension
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def num_chars(self):
|
|||
|
|
return len(self._vocab)
|
|||
|
|
|
|||
|
|
def _create_vocab(self):
|
|||
|
|
_vocab = self._characters
|
|||
|
|
if self.is_unique:
|
|||
|
|
_vocab = list(set(_vocab))
|
|||
|
|
if self.is_sorted:
|
|||
|
|
_vocab = sorted(_vocab)
|
|||
|
|
_vocab = list(_vocab)
|
|||
|
|
_vocab = [self._blank] + _vocab if self._blank is not None and len(self._blank) > 0 else _vocab
|
|||
|
|
_vocab = [self._bos] + _vocab if self._bos is not None and len(self._bos) > 0 else _vocab
|
|||
|
|
_vocab = [self._eos] + _vocab if self._eos is not None and len(self._eos) > 0 else _vocab
|
|||
|
|
_vocab = [self._pad] + _vocab if self._pad is not None and len(self._pad) > 0 else _vocab
|
|||
|
|
self.vocab = _vocab + list(self._punctuations)
|
|||
|
|
if self.is_unique:
|
|||
|
|
duplicates = {x for x in self.vocab if self.vocab.count(x) > 1}
|
|||
|
|
assert (
|
|||
|
|
len(self.vocab) == len(self._char_to_id) == len(self._id_to_char)
|
|||
|
|
), f" [!] There are duplicate characters in the character set. {duplicates}"
|
|||
|
|
|
|||
|
|
def char_to_id(self, char: str) -> int:
|
|||
|
|
try:
|
|||
|
|
return self._char_to_id[char]
|
|||
|
|
except KeyError as e:
|
|||
|
|
raise KeyError(f" [!] {repr(char)} is not in the vocabulary.") from e
|
|||
|
|
|
|||
|
|
def id_to_char(self, idx: int) -> str:
|
|||
|
|
return self._id_to_char[idx]
|
|||
|
|
|
|||
|
|
def print_log(self, level: int = 0):
|
|||
|
|
"""
|
|||
|
|
Prints the vocabulary in a nice format.
|
|||
|
|
"""
|
|||
|
|
indent = "\t" * level
|
|||
|
|
print(f"{indent}| > Characters: {self._characters}")
|
|||
|
|
print(f"{indent}| > Punctuations: {self._punctuations}")
|
|||
|
|
print(f"{indent}| > Pad: {self._pad}")
|
|||
|
|
print(f"{indent}| > EOS: {self._eos}")
|
|||
|
|
print(f"{indent}| > BOS: {self._bos}")
|
|||
|
|
print(f"{indent}| > Blank: {self._blank}")
|
|||
|
|
print(f"{indent}| > Vocab: {self.vocab}")
|
|||
|
|
print(f"{indent}| > Num chars: {self.num_chars}")
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def init_from_config(config: "Coqpit"): # pylint: disable=unused-argument
|
|||
|
|
"""Init your character class from a config.
|
|||
|
|
|
|||
|
|
Implement this method for your subclass.
|
|||
|
|
"""
|
|||
|
|
# use character set from config
|
|||
|
|
if config.characters is not None:
|
|||
|
|
return BaseCharacters(**config.characters), config
|
|||
|
|
# return default character set
|
|||
|
|
characters = BaseCharacters()
|
|||
|
|
new_config = replace(config, characters=characters.to_config())
|
|||
|
|
return characters, new_config
|
|||
|
|
|
|||
|
|
def to_config(self) -> "CharactersConfig":
|
|||
|
|
return CharactersConfig(
|
|||
|
|
characters=self._characters,
|
|||
|
|
punctuations=self._punctuations,
|
|||
|
|
pad=self._pad,
|
|||
|
|
eos=self._eos,
|
|||
|
|
bos=self._bos,
|
|||
|
|
blank=self._blank,
|
|||
|
|
is_unique=self.is_unique,
|
|||
|
|
is_sorted=self.is_sorted,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class IPAPhonemes(BaseCharacters):
|
|||
|
|
"""🐸IPAPhonemes class to manage `TTS.tts` model vocabulary
|
|||
|
|
|
|||
|
|
Intended to be used with models using IPAPhonemes as input.
|
|||
|
|
It uses system defaults for the undefined class arguments.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
characters (str):
|
|||
|
|
Main set of case-sensitive characters to be used in the vocabulary. Defaults to `_phonemes`.
|
|||
|
|
|
|||
|
|
punctuations (str):
|
|||
|
|
Characters to be treated as punctuation. Defaults to `_punctuations`.
|
|||
|
|
|
|||
|
|
pad (str):
|
|||
|
|
Special padding character that would be ignored by the model. Defaults to `_pad`.
|
|||
|
|
|
|||
|
|
eos (str):
|
|||
|
|
End of the sentence character. Defaults to `_eos`.
|
|||
|
|
|
|||
|
|
bos (str):
|
|||
|
|
Beginning of the sentence character. Defaults to `_bos`.
|
|||
|
|
|
|||
|
|
blank (str):
|
|||
|
|
Optional character used between characters by some models for better prosody. Defaults to `_blank`.
|
|||
|
|
|
|||
|
|
is_unique (bool):
|
|||
|
|
Remove duplicates from the provided characters. Defaults to True.
|
|||
|
|
|
|||
|
|
is_sorted (bool):
|
|||
|
|
Sort the characters in alphabetical order. Defaults to True.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
characters: str = _phonemes,
|
|||
|
|
punctuations: str = _punctuations,
|
|||
|
|
pad: str = _pad,
|
|||
|
|
eos: str = _eos,
|
|||
|
|
bos: str = _bos,
|
|||
|
|
blank: str = _blank,
|
|||
|
|
is_unique: bool = False,
|
|||
|
|
is_sorted: bool = True,
|
|||
|
|
) -> None:
|
|||
|
|
super().__init__(characters, punctuations, pad, eos, bos, blank, is_unique, is_sorted)
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def init_from_config(config: "Coqpit"):
|
|||
|
|
"""Init a IPAPhonemes object from a model config
|
|||
|
|
|
|||
|
|
If characters are not defined in the config, it will be set to the default characters and the config
|
|||
|
|
will be updated.
|
|||
|
|
"""
|
|||
|
|
# band-aid for compatibility with old models
|
|||
|
|
if "characters" in config and config.characters is not None:
|
|||
|
|
if "phonemes" in config.characters and config.characters.phonemes is not None:
|
|||
|
|
config.characters["characters"] = config.characters["phonemes"]
|
|||
|
|
return (
|
|||
|
|
IPAPhonemes(
|
|||
|
|
characters=config.characters["characters"],
|
|||
|
|
punctuations=config.characters["punctuations"],
|
|||
|
|
pad=config.characters["pad"],
|
|||
|
|
eos=config.characters["eos"],
|
|||
|
|
bos=config.characters["bos"],
|
|||
|
|
blank=config.characters["blank"],
|
|||
|
|
is_unique=config.characters["is_unique"],
|
|||
|
|
is_sorted=config.characters["is_sorted"],
|
|||
|
|
),
|
|||
|
|
config,
|
|||
|
|
)
|
|||
|
|
# use character set from config
|
|||
|
|
if config.characters is not None:
|
|||
|
|
return IPAPhonemes(**config.characters), config
|
|||
|
|
# return default character set
|
|||
|
|
characters = IPAPhonemes()
|
|||
|
|
new_config = replace(config, characters=characters.to_config())
|
|||
|
|
return characters, new_config
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Graphemes(BaseCharacters):
|
|||
|
|
"""🐸Graphemes class to manage `TTS.tts` model vocabulary
|
|||
|
|
|
|||
|
|
Intended to be used with models using graphemes as input.
|
|||
|
|
It uses system defaults for the undefined class arguments.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
characters (str):
|
|||
|
|
Main set of case-sensitive characters to be used in the vocabulary. Defaults to `_characters`.
|
|||
|
|
|
|||
|
|
punctuations (str):
|
|||
|
|
Characters to be treated as punctuation. Defaults to `_punctuations`.
|
|||
|
|
|
|||
|
|
pad (str):
|
|||
|
|
Special padding character that would be ignored by the model. Defaults to `_pad`.
|
|||
|
|
|
|||
|
|
eos (str):
|
|||
|
|
End of the sentence character. Defaults to `_eos`.
|
|||
|
|
|
|||
|
|
bos (str):
|
|||
|
|
Beginning of the sentence character. Defaults to `_bos`.
|
|||
|
|
|
|||
|
|
is_unique (bool):
|
|||
|
|
Remove duplicates from the provided characters. Defaults to True.
|
|||
|
|
|
|||
|
|
is_sorted (bool):
|
|||
|
|
Sort the characters in alphabetical order. Defaults to True.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
characters: str = _characters,
|
|||
|
|
punctuations: str = _punctuations,
|
|||
|
|
pad: str = _pad,
|
|||
|
|
eos: str = _eos,
|
|||
|
|
bos: str = _bos,
|
|||
|
|
blank: str = _blank,
|
|||
|
|
is_unique: bool = False,
|
|||
|
|
is_sorted: bool = True,
|
|||
|
|
) -> None:
|
|||
|
|
super().__init__(characters, punctuations, pad, eos, bos, blank, is_unique, is_sorted)
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def init_from_config(config: "Coqpit"):
|
|||
|
|
"""Init a Graphemes object from a model config
|
|||
|
|
|
|||
|
|
If characters are not defined in the config, it will be set to the default characters and the config
|
|||
|
|
will be updated.
|
|||
|
|
"""
|
|||
|
|
if config.characters is not None:
|
|||
|
|
# band-aid for compatibility with old models
|
|||
|
|
if "phonemes" in config.characters:
|
|||
|
|
return (
|
|||
|
|
Graphemes(
|
|||
|
|
characters=config.characters["characters"],
|
|||
|
|
punctuations=config.characters["punctuations"],
|
|||
|
|
pad=config.characters["pad"],
|
|||
|
|
eos=config.characters["eos"],
|
|||
|
|
bos=config.characters["bos"],
|
|||
|
|
blank=config.characters["blank"],
|
|||
|
|
is_unique=config.characters["is_unique"],
|
|||
|
|
is_sorted=config.characters["is_sorted"],
|
|||
|
|
),
|
|||
|
|
config,
|
|||
|
|
)
|
|||
|
|
return Graphemes(**config.characters), config
|
|||
|
|
characters = Graphemes()
|
|||
|
|
new_config = replace(config, characters=characters.to_config())
|
|||
|
|
return characters, new_config
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
gr = Graphemes()
|
|||
|
|
ph = IPAPhonemes()
|
|||
|
|
gr.print_log()
|
|||
|
|
ph.print_log()
|