mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-22 08:16:08 +00:00
Brandon Thomas
353 lines
11 KiB
Python
353 lines
11 KiB
Python
# Copyright (c) 2020, NVIDIA CORPORATION. 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 abc
|
|
import itertools
|
|
import re
|
|
import string
|
|
import unicodedata
|
|
from builtins import str as unicode
|
|
from typing import List
|
|
|
|
import nltk
|
|
|
|
from nemo.collections.common.parts.preprocessing import parsers
|
|
|
|
try:
|
|
import g2p_en # noqa
|
|
|
|
_g2p = g2p_en.G2p()
|
|
_g2p.variables = None
|
|
|
|
HAVE_G2P = True
|
|
except (FileNotFoundError, LookupError):
|
|
HAVE_G2P = False
|
|
|
|
_words_re = re.compile("([a-z\-]+'[a-z\-]+|[a-z\-]+)|([^a-z{}]+)")
|
|
|
|
|
|
def _text_preprocessing(text):
|
|
text = unicode(text)
|
|
text = ''.join(char for char in unicodedata.normalize('NFD', text) if unicodedata.category(char) != 'Mn')
|
|
text = text.lower()
|
|
text = re.sub("[^ a-z'\".,?!()\[\]:;\-]", "", text)
|
|
return text
|
|
|
|
|
|
def _word_tokenize(text):
|
|
words = _words_re.findall(text)
|
|
words = [re.sub(r'\s(\d)', r'\1', word[1].upper()) if word[0] == '' else word[0] for word in words]
|
|
return words
|
|
|
|
|
|
class G2p:
|
|
def __init__(
|
|
self,
|
|
phoneme_dict_path=None,
|
|
use_seq2seq_for_oov=False,
|
|
ignore_ambiguous_words=True,
|
|
text_preprocessing_func=_text_preprocessing,
|
|
word_tokenize_func=_word_tokenize,
|
|
):
|
|
# Download NLTK datasets if this class is to be instantiated
|
|
try:
|
|
nltk.data.find('taggers/averaged_perceptron_tagger.zip')
|
|
except LookupError:
|
|
nltk.download('averaged_perceptron_tagger', quiet=True)
|
|
try:
|
|
nltk.data.find('corpora/cmudict.zip')
|
|
except LookupError:
|
|
nltk.download('cmudict', quiet=True)
|
|
|
|
self.homograph2features = _g2p.homograph2features
|
|
self.g2p_dict = self._construct_grapheme2phoneme_dict(phoneme_dict_path)
|
|
self.use_seq2seq_for_oov = use_seq2seq_for_oov
|
|
self.ignore_ambiguous_words = ignore_ambiguous_words
|
|
|
|
self.text_preprocessing_func = text_preprocessing_func
|
|
self.word_tokenize_func = word_tokenize_func
|
|
|
|
@staticmethod
|
|
def _construct_grapheme2phoneme_dict(phoneme_dict_path=None, encoding='latin-1'):
|
|
if phoneme_dict_path is None:
|
|
from nltk.corpus import cmudict
|
|
|
|
return cmudict.dict()
|
|
|
|
_alt_re = re.compile(r'\([0-9]+\)')
|
|
g2p_dict = {}
|
|
with open(phoneme_dict_path, encoding=encoding) as file:
|
|
for line in file:
|
|
if len(line) and ('A' <= line[0] <= 'Z' or line[0] == "'"):
|
|
parts = line.split(' ')
|
|
word = re.sub(_alt_re, '', parts[0])
|
|
word = word.lower()
|
|
|
|
pronunciation = parts[1].strip().split(" ")
|
|
if word in g2p_dict:
|
|
g2p_dict[word].append(pronunciation)
|
|
else:
|
|
g2p_dict[word] = [pronunciation]
|
|
return g2p_dict
|
|
|
|
def handle_ambiguous(self, word):
|
|
if not self.ignore_ambiguous_words or len(self.g2p_dict[word]) == 1:
|
|
return True
|
|
return False
|
|
|
|
def __call__(self, text):
|
|
text = self.text_preprocessing_func(text)
|
|
words = self.word_tokenize_func(text)
|
|
words_and_pos_tags = nltk.pos_tag(words)
|
|
|
|
prons = []
|
|
for word, pos in words_and_pos_tags:
|
|
word_by_hyphen = word.split("-")
|
|
|
|
# punctuation
|
|
if re.search("[a-zA-Z]", word) is None:
|
|
pron = list(word)
|
|
# homograph
|
|
elif word in self.homograph2features:
|
|
pron1, pron2, pos1 = self.homograph2features[word]
|
|
if pos.startswith(pos1):
|
|
pron = pron1
|
|
else:
|
|
pron = pron2
|
|
# `'s` suffix
|
|
elif (
|
|
len(word) > 2
|
|
and word.endswith("'s")
|
|
and (word not in self.g2p_dict)
|
|
and (word[:-2] in self.g2p_dict)
|
|
and self.handle_ambiguous(word[:-2])
|
|
):
|
|
pron = self.g2p_dict[word[:-2]][0] + ["Z"]
|
|
# `s` suffix
|
|
elif (
|
|
len(word) > 1
|
|
and word.endswith("s")
|
|
and (word not in self.g2p_dict)
|
|
and (word[:-1] in self.g2p_dict)
|
|
and self.handle_ambiguous(word[:-1])
|
|
):
|
|
pron = self.g2p_dict[word[:-1]][0] + ["Z"]
|
|
# g2p dict
|
|
elif word in self.g2p_dict and self.handle_ambiguous(word):
|
|
pron = self.g2p_dict[word][0]
|
|
# word with hyphens
|
|
elif len(word_by_hyphen) > 1 and all(
|
|
[sub_word in self.g2p_dict and self.handle_ambiguous(sub_word) for sub_word in word_by_hyphen]
|
|
):
|
|
pron = []
|
|
for sub_word in word_by_hyphen:
|
|
pron.extend(self.g2p_dict[sub_word][0])
|
|
pron.extend(["-"])
|
|
pron.pop()
|
|
else:
|
|
if self.use_seq2seq_for_oov:
|
|
# run gru-based seq2seq model from _g2p
|
|
pron = _g2p.predict(word)
|
|
else:
|
|
pron = word
|
|
|
|
prons.extend(pron)
|
|
|
|
return prons
|
|
|
|
|
|
class Base(abc.ABC):
|
|
"""Vocabulary for turning str text to list of int tokens."""
|
|
|
|
# fmt: off
|
|
PUNCT = ( # Derived from LJSpeech and "/" additionally
|
|
',', '.', '!', '?', '-',
|
|
':', ';', '/', '"', '(',
|
|
')', '[', ']', '{', '}',
|
|
)
|
|
# fmt: on
|
|
PAD, BLANK, OOV = '<pad>', '<blank>', '<oov>'
|
|
|
|
def __init__(self, labels, *, pad=PAD, blank=BLANK, oov=OOV, sep='', add_blank_at="last_but_one"):
|
|
super().__init__()
|
|
|
|
labels = list(labels)
|
|
self.pad, labels = len(labels), labels + [pad] # Padding
|
|
|
|
if add_blank_at is not None:
|
|
self.blank, labels = len(labels), labels + [blank] # Reserved for blank from QN
|
|
else:
|
|
# use add_blank_at=None only for ASR where blank is added automatically
|
|
self.blank = -1
|
|
|
|
self.oov, labels = len(labels), labels + [oov] # Out Of Vocabulary
|
|
|
|
if add_blank_at == "last":
|
|
labels[-1], labels[-2] = labels[-2], labels[-1]
|
|
self.oov, self.blank = self.blank, self.oov
|
|
|
|
self.labels = labels
|
|
self.sep = sep
|
|
|
|
self._util_ids = {self.pad, self.blank, self.oov}
|
|
self._label2id = {l: i for i, l in enumerate(labels)}
|
|
self._id2label = labels
|
|
|
|
def __call__(self, text: str) -> List[int]:
|
|
return self.encode(text)
|
|
|
|
@abc.abstractmethod
|
|
def encode(self, text: str) -> List[int]:
|
|
"""Turns str text into int tokens."""
|
|
pass
|
|
|
|
def decode(self, tokens: List[int]) -> str:
|
|
"""Turns ints tokens into str text."""
|
|
return self.sep.join(self._id2label[t] for t in tokens if t not in self._util_ids)
|
|
|
|
|
|
class Chars(Base):
|
|
"""Chars vocabulary."""
|
|
|
|
def __init__(
|
|
self, punct=True, spaces=False, apostrophe=True, add_blank_at="last_but_one",
|
|
):
|
|
labels = []
|
|
self.space, labels = len(labels), labels + [' '] # Space
|
|
labels.extend(string.ascii_lowercase)
|
|
if apostrophe:
|
|
labels.append("'") # Apostrophe for saving "don't" and "Joe's"
|
|
|
|
if punct:
|
|
labels.extend(self.PUNCT)
|
|
|
|
super().__init__(labels, add_blank_at=add_blank_at)
|
|
|
|
self.punct = punct
|
|
self.spaces = spaces
|
|
|
|
self._parser = parsers.ENCharParser(labels)
|
|
|
|
def encode(self, text):
|
|
"""See base class."""
|
|
text = self._parser._normalize(text) # noqa
|
|
|
|
if self.spaces:
|
|
for p in set(text) & set(self.PUNCT):
|
|
text = text.replace(p, f' {p} ')
|
|
text = text.strip().replace(' ', ' ')
|
|
|
|
return self._parser._tokenize(text) # noqa
|
|
|
|
|
|
class Phonemes(Base):
|
|
"""Phonemes vocabulary."""
|
|
|
|
# fmt: off
|
|
VOWELS = (
|
|
'AA', 'AE', 'AH', 'AO', 'AW',
|
|
'AY', 'EH', 'ER', 'EY', 'IH',
|
|
'IY', 'OW', 'OY', 'UH', 'UW',
|
|
)
|
|
CONSONANTS = (
|
|
'B', 'CH', 'D', 'DH', 'F', 'G',
|
|
'HH', 'JH', 'K', 'L', 'M', 'N',
|
|
'NG', 'P', 'R', 'S', 'SH', 'T',
|
|
'TH', 'V', 'W', 'Y', 'Z', 'ZH',
|
|
)
|
|
# fmt: on
|
|
|
|
def __init__(
|
|
self,
|
|
punct=True,
|
|
stresses=False,
|
|
spaces=True,
|
|
chars=False,
|
|
*,
|
|
space=' ',
|
|
silence=None,
|
|
apostrophe=True,
|
|
oov=Base.OOV,
|
|
sep='|', # To be able to distinguish between 2/3 letters codes.
|
|
add_blank_at="last_but_one",
|
|
pad_with_space=False,
|
|
improved_version_g2p=False,
|
|
phoneme_dict_path=None,
|
|
):
|
|
labels = []
|
|
self.space, labels = len(labels), labels + [space] # Space
|
|
|
|
if silence is not None:
|
|
self.silence, labels = len(labels), labels + [silence] # Silence
|
|
|
|
labels.extend(self.CONSONANTS)
|
|
vowels = list(self.VOWELS)
|
|
|
|
if stresses:
|
|
vowels = [f'{p}{s}' for p, s in itertools.product(vowels, (0, 1, 2))]
|
|
labels.extend(vowels)
|
|
|
|
if chars:
|
|
labels.extend(string.ascii_lowercase)
|
|
|
|
if apostrophe:
|
|
labels.append("'") # Apostrophe
|
|
|
|
if punct:
|
|
labels.extend(self.PUNCT)
|
|
|
|
super().__init__(labels, oov=oov, sep=sep, add_blank_at=add_blank_at)
|
|
|
|
self.punct = punct
|
|
self.stresses = stresses
|
|
self.spaces = spaces
|
|
self.pad_with_space = pad_with_space
|
|
|
|
if improved_version_g2p:
|
|
self.g2p = G2p(phoneme_dict_path)
|
|
else:
|
|
self.g2p = _g2p
|
|
|
|
def encode(self, text):
|
|
"""See base class."""
|
|
ps, space, labels = [], self.labels[self.space], set(self.labels)
|
|
|
|
for p in self.g2p(text): # noqa
|
|
# Remove stress
|
|
if p.isalnum() and len(p) == 3 and not self.stresses:
|
|
p = p[:2]
|
|
|
|
# Add space if last one isn't one
|
|
if p == space and len(ps) > 0 and ps[-1] != space:
|
|
ps.append(p)
|
|
|
|
# Add next phoneme
|
|
if (p.isalnum() or p == "'") and p in labels:
|
|
ps.append(p)
|
|
|
|
# Add punct and remove space if needed
|
|
if (p in self.PUNCT) and self.punct:
|
|
if not self.spaces and len(ps) > 0 and ps[-1] == space:
|
|
ps.pop()
|
|
ps.append(p)
|
|
|
|
# Remove trailing spaces
|
|
while ps[-1] == space:
|
|
ps.pop()
|
|
|
|
if self.pad_with_space:
|
|
ps = [space] + ps + [space]
|
|
|
|
return [self._label2id[p] for p in ps]
|