Spaces:
Runtime error
Runtime error
File size: 4,843 Bytes
2045faa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
import os, sys
import tensorflow as tf
import numpy as np
from tensorflow.keras.models import Model
from tensorflow.keras import layers
import datetime
sys.path.append(os.path.dirname(__file__))
import encoder, extractor, discriminator, log_melspectrogram, speech_embedding
from utils import make_feature_matrix as concat_sequence
seed = 42
tf.random.set_seed(seed)
np.random.seed(seed)
class ukws(Model):
"""Base class for user-defined kws mdoel"""
def __init__(self, name="ukws", **kwargs):
super(ukws, self).__init__(name=name)
def call(self, speech, text):
"""
Args:
speech : speech feature of shape `(batch, time)`
text : text embedding of shape `(batch, phoneme)`
"""
raise NotImplementedError
class BaseUKWS(ukws):
"""Base class for user-defined kws mdoel"""
def __init__(self, name="BaseUKWS", **kwargs):
super(BaseUKWS, self).__init__(name=name)
embedding=128
self.audio_input = kwargs['audio_input']
self.text_input = kwargs['text_input']
self.stack_extractor = kwargs['stack_extractor']
_stft={
'frame_length' : kwargs['frame_length'],
'hop_length' : kwargs['hop_length'],
'num_mel' : kwargs['num_mel'] ,
'sample_rate' : kwargs['sample_rate'],
'log_mel' : kwargs['log_mel'],
}
_ae = {
# [filter, kernel size, stride]
'conv' : [[embedding, 5, 2], [embedding * 2, 5, 1]],
# [unit]
'gru' : [[embedding], [embedding]],
# fully-connected layer unit
'fc' : embedding,
'audio_input' : self.audio_input,
}
_te = {
# fully-connected layer unit
'fc' : embedding,
# number of uniq. phonemes
'vocab' : kwargs['vocab'],
'text_input' : kwargs['text_input'],
}
_ext = {
# [unit]
'embedding' : embedding,
}
_dis = {
# [unit]
'gru' : [[embedding],],
}
if self.audio_input == 'both':
self.SPEC = log_melspectrogram.LogMelgramLayer(**_stft)
self.EMBD = speech_embedding.GoogleSpeechEmbedder()
self.AE = encoder.EfficientAudioEncoder(downsample=False, **_ae)
else:
if self.audio_input == 'raw':
self.FEAT = log_melspectrogram.LogMelgramLayer(**_stft)
elif self.audio_input == 'google_embed':
self.FEAT = speech_embedding.GoogleSpeechEmbedder()
self.AE = encoder.AudioEncoder(**_ae)
self.TE = encoder.TextEncoder(**_te)
if kwargs['stack_extractor']:
self.EXT = extractor.StackExtractor(**_ext)
else:
self.EXT = extractor.BaseExtractor(**_ext)
self.DIS = discriminator.BaseDiscriminator(**_dis)
self.seq_ce_logit = layers.Dense(1, name='sequence_ce')
def call(self, speech, text):
"""
Args:
speech : speech feature of shape `(batch, time)`
text : text embedding of shape `(batch, phoneme)`
"""
if self.audio_input == 'both':
s = self.SPEC(speech)
g = self.EMBD(speech)
emb_s, LDN = self.AE(s, g)
else:
feat = self.FEAT(speech)
emb_s, LDN = self.AE(feat)
emb_t = self.TE(text)
attention_output, affinity_matrix = self.EXT(emb_s, emb_t)
prob, LD = self.DIS(attention_output)
n_speech = tf.math.reduce_sum(tf.cast(emb_s._keras_mask, tf.float32), -1)
if self.stack_extractor:
n_speech = tf.math.reduce_sum(tf.cast(emb_s._keras_mask, tf.float32), -1)
n_text = tf.math.reduce_sum(tf.cast(emb_t._keras_mask, tf.float32), -1)
n_total = n_speech + n_text
valid_mask = tf.sequence_mask(n_total, maxlen=tf.shape(attention_output)[1], dtype=tf.float32) - tf.sequence_mask(n_speech, maxlen=tf.shape(attention_output)[1], dtype=tf.float32)
valid_attention_output = tf.ragged.boolean_mask(attention_output, tf.cast(valid_mask, tf.bool)).to_tensor(0.)
seq_ce_logit = self.seq_ce_logit(valid_attention_output)[:,:,0]
seq_ce_logit = tf.pad(seq_ce_logit, [[0, 0],[0, tf.shape(emb_t)[1] - tf.shape(seq_ce_logit)[1]]], 'CONSTANT', constant_values=0.)
seq_ce_logit._keras_mask = emb_t._keras_mask
else:
seq_ce_logit = self.seq_ce_logit(attention_output)[:,:,0]
seq_ce_logit._keras_mask = attention_output._keras_mask
return prob, affinity_matrix, LD, seq_ce_logit, n_speech |