Move data to device and compute text seq . length

def fetch_data(self, data):
        ''' Move data to device and compute text seq. length'''
        _, feat, feat_len, txt = data
        feat = feat.to(self.device)
        feat_len = feat_len.to(self.device)
        txt = txt.to(self.device)
        txt_len = torch.sum(txt != 0, dim=-1)

        return feat, feat_len, txt, txt_len

def fetch_data ( self , data ) : _ , feat , feat_len , txt = data feat = feat . to ( self . device ) feat_len = feat_len . to ( self . device ) txt = txt . to ( self . device ) txt_len = torch . sum ( txt != 0 , dim = - 1 ) return feat , feat_len , txt , txt_len


Load data for training / validation store tokenizer and input / output shape

def load_data(self):
        ''' Load data for training/validation, store tokenizer and input/output shape'''
        self.tr_set, self.dv_set, self.feat_dim, self.vocab_size, self.tokenizer, msg = \
            load_dataset(self.paras.njobs, self.paras.gpu, self.paras.pin_memory,
                         self.curriculum > 0, **self.config['data'])
        self.verbose(msg)

def load_data ( self ) : self . tr_set , self . dv_set , self . feat_dim , self . vocab_size , self . tokenizer , msg = load_dataset ( self . paras . njobs , self . paras . gpu , self . paras . pin_memory , self . curriculum > 0 , * * self . config [ 'data' ] ) self . verbose ( msg )


def set_model(self):
        ''' Setup ASR model and optimizer '''
        # Model
        init_adadelta = self.config['hparas']['optimizer'] == 'Adadelta'
        self.model = ASR(self.feat_dim, self.vocab_size, init_adadelta, **
                         self.config['model']).to(self.device)
        self.verbose(self.model.create_msg())
        model_paras = [{'params': self.model.parameters()}]

        # Losses
        self.seq_loss = torch.nn.CrossEntropyLoss(ignore_index=0)
        # Note: zero_infinity=False is unstable?
        self.ctc_loss = torch.nn.CTCLoss(blank=0, zero_infinity=False)

        # Plug-ins
        self.emb_fuse = False
        self.emb_reg = ('emb' in self.config) and (
            self.config['emb']['enable'])
        if self.emb_reg:
            from src.plugin import EmbeddingRegularizer
            self.emb_decoder = EmbeddingRegularizer(
                self.tokenizer, self.model.dec_dim, **self.config['emb']).to(self.device)
            model_paras.append({'params': self.emb_decoder.parameters()})
            self.emb_fuse = self.emb_decoder.apply_fuse
            if self.emb_fuse:
                self.seq_loss = torch.nn.NLLLoss(ignore_index=0)
            self.verbose(self.emb_decoder.create_msg())

        # Optimizer
        self.optimizer = Optimizer(model_paras, **self.config['hparas'])
        self.verbose(self.optimizer.create_msg())

        # Enable AMP if needed
        self.enable_apex()

        # Automatically load pre-trained model if self.paras.load is given
        self.load_ckpt()


def exec(self):
        ''' Training End-to-end ASR system '''
        self.verbose('Total training steps {}.'.format(
            human_format(self.max_step)))
        ctc_loss, att_loss, emb_loss = None, None, None
        n_epochs = 0
        self.timer.set()

        while self.step < self.max_step:
            # Renew dataloader to enable random sampling
            if self.curriculum > 0 and n_epochs == self.curriculum:
                self.verbose(
                    'Curriculum learning ends after {} epochs, starting random sampling.'.format(n_epochs))
                self.tr_set, _, _, _, _, _ = \
                    load_dataset(self.paras.njobs, self.paras.gpu, self.paras.pin_memory,
                                 False, **self.config['data'])
            for data in self.tr_set:
                # Pre-step : update tf_rate/lr_rate and do zero_grad
                tf_rate = self.optimizer.pre_step(self.step)
                total_loss = 0

                # Fetch data
                feat, feat_len, txt, txt_len = self.fetch_data(data)
                self.timer.cnt('rd')

                # Forward model
                # Note: txt should NOT start w/ <sos>
                ctc_output, encode_len, att_output, att_align, dec_state = \
                    self.model(feat, feat_len, max(txt_len), tf_rate=tf_rate,
                               teacher=txt, get_dec_state=self.emb_reg)

                # Plugins
                if self.emb_reg:
                    emb_loss, fuse_output = self.emb_decoder(
                        dec_state, att_output, label=txt)
                    total_loss += self.emb_decoder.weight*emb_loss

                # Compute all objectives
                if ctc_output is not None:
                    if self.paras.cudnn_ctc:
                        ctc_loss = self.ctc_loss(ctc_output.transpose(0, 1),
                                                 txt.to_sparse().values().to(device='cpu', dtype=torch.int32),
                                                 [ctc_output.shape[1]] *
                                                 len(ctc_output),
                                                 txt_len.cpu().tolist())
                    else:
                        ctc_loss = self.ctc_loss(ctc_output.transpose(
                            0, 1), txt, encode_len, txt_len)
                    total_loss += ctc_loss*self.model.ctc_weight

                if att_output is not None:
                    b, t, _ = att_output.shape
                    att_output = fuse_output if self.emb_fuse else att_output
                    att_loss = self.seq_loss(
                        att_output.view(b*t, -1), txt.view(-1))
                    total_loss += att_loss*(1-self.model.ctc_weight)

                self.timer.cnt('fw')

                # Backprop
                grad_norm = self.backward(total_loss)
                self.step += 1

                # Logger
                if (self.step == 1) or (self.step % self.PROGRESS_STEP == 0):
                    self.progress('Tr stat | Loss - {:.2f} | Grad. Norm - {:.2f} | {}'
                                  .format(total_loss.cpu().item(), grad_norm, self.timer.show()))
                    self.write_log(
                        'loss', {'tr_ctc': ctc_loss, 'tr_att': att_loss})
                    self.write_log('emb_loss', {'tr': emb_loss})
                    self.write_log('wer', {'tr_att': cal_er(self.tokenizer, att_output, txt),
                                           'tr_ctc': cal_er(self.tokenizer, ctc_output, txt, ctc=True)})
                    if self.emb_fuse:
                        if self.emb_decoder.fuse_learnable:
                            self.write_log('fuse_lambda', {
                                           'emb': self.emb_decoder.get_weight()})
                        self.write_log(
                            'fuse_temp', {'temp': self.emb_decoder.get_temp()})

                # Validation
                if (self.step == 1) or (self.step % self.valid_step == 0):
                    self.validate()

                # End of step
                # https://github.com/pytorch/pytorch/issues/13246#issuecomment-529185354
                torch.cuda.empty_cache()
                self.timer.set()
                if self.step > self.max_step:
                    break
            n_epochs += 1
        self.log.close()


Move data to device insert <sos > and compute text seq . length

def fetch_data(self, data):
        ''' Move data to device, insert <sos> and compute text seq. length'''
        txt = torch.cat(
            (torch.zeros((data.shape[0], 1), dtype=torch.long), data), dim=1).to(self.device)
        txt_len = torch.sum(data != 0, dim=-1)
        return txt, txt_len

def fetch_data ( self , data ) : txt = torch . cat ( ( torch . zeros ( ( data . shape [ 0 ] , 1 ) , dtype = torch . long ) , data ) , dim = 1 ) . to ( self . device ) txt_len = torch . sum ( data != 0 , dim = - 1 ) return txt , txt_len


def load_data(self):
        ''' Load data for training/validation, store tokenizer and input/output shape'''
        self.tr_set, self.dv_set, self.vocab_size, self.tokenizer, msg = \
            load_textset(self.paras.njobs, self.paras.gpu,
                         self.paras.pin_memory, **self.config['data'])
        self.verbose(msg)

def load_data ( self ) : self . tr_set , self . dv_set , self . vocab_size , self . tokenizer , msg = load_textset ( self . paras . njobs , self . paras . gpu , self . paras . pin_memory , * * self . config [ 'data' ] ) self . verbose ( msg )


def set_model(self):
        ''' Setup ASR model and optimizer '''

        # Model
        self.model = RNNLM(self.vocab_size, **
                           self.config['model']).to(self.device)
        self.verbose(self.model.create_msg())
        # Losses
        self.seq_loss = torch.nn.CrossEntropyLoss(ignore_index=0)
        # Optimizer
        self.optimizer = Optimizer(
            self.model.parameters(), **self.config['hparas'])
        # Enable AMP if needed
        self.enable_apex()
        # load pre-trained model
        if self.paras.load:
            self.load_ckpt()
            ckpt = torch.load(self.paras.load, map_location=self.device)
            self.model.load_state_dict(ckpt['model'])
            self.optimizer.load_opt_state_dict(ckpt['optimizer'])
            self.step = ckpt['global_step']
            self.verbose('Load ckpt from {}, restarting at step {}'.format(
                self.paras.load, self.step))


def exec(self):
        ''' Training End-to-end ASR system '''
        self.verbose('Total training steps {}.'.format(
            human_format(self.max_step)))
        self.timer.set()

        while self.step < self.max_step:
            for data in self.tr_set:
                # Pre-step : update tf_rate/lr_rate and do zero_grad
                self.optimizer.pre_step(self.step)

                # Fetch data
                txt, txt_len = self.fetch_data(data)
                self.timer.cnt('rd')

                # Forward model
                pred, _ = self.model(txt[:, :-1], txt_len)

                # Compute all objectives
                lm_loss = self.seq_loss(
                    pred.view(-1, self.vocab_size), txt[:, 1:].reshape(-1))
                self.timer.cnt('fw')

                # Backprop
                grad_norm = self.backward(lm_loss)
                self.step += 1

                # Logger
                if self.step % self.PROGRESS_STEP == 0:
                    self.progress('Tr stat | Loss - {:.2f} | Grad. Norm - {:.2f} | {}'
                                  .format(lm_loss.cpu().item(), grad_norm, self.timer.show()))
                    self.write_log('entropy', {'tr': lm_loss})
                    self.write_log(
                        'perplexity', {'tr': torch.exp(lm_loss).cpu().item()})

                # Validation
                if (self.step == 1) or (self.step % self.valid_step == 0):
                    self.validate()

                # End of step
                self.timer.set()
                if self.step > self.max_step:
                    break
        self.log.close()


Get transcription of target wave file it s somewhat redundant for accessing each txt multiplt times but it works fine with multi - thread

def read_text(file):
    '''Get transcription of target wave file,
       it's somewhat redundant for accessing each txt multiplt times,
       but it works fine with multi-thread'''
    src_file = '-'.join(file.split('-')[:-1])+'.trans.txt'
    idx = file.split('/')[-1].split('.')[0]

    with open(src_file, 'r') as fp:
        for line in fp:
            if idx == line.split(' ')[0]:
                return line[:-1].split(' ', 1)[1]

def
read_text
(
file
)
:
src_file
=
'-'
.
join
(
file
.
split
(
'-'
)
[
:
1
]
)
+
'.trans.txt'
idx
=
file
.
split
(
'/'
)
[
1
]
.
split
(
'.'
)
[
0
]
with
open
(
src_file
,
'r'
)
as
fp
:
for
line
in
fp
:
if
idx
==
line
.
split
(
' '
)
[
0
]
:
return
line
[
:
1
]
.
split
(
' '
,
1
)
[
1
]


Generate bert s embedding from fine - tuned model .

def generate_embedding(bert_model, labels):
    """Generate bert's embedding from fine-tuned model."""
    batch_size, time = labels.shape

    cls_ids = torch.full(
        (batch_size, 1), bert_model.bert_text_encoder.cls_idx, dtype=labels.dtype, device=labels.device)
    bert_labels = torch.cat([cls_ids, labels], 1)
    # replace eos with sep
    eos_idx = bert_model.bert_text_encoder.eos_idx
    sep_idx = bert_model.bert_text_encoder.sep_idx
    bert_labels[bert_labels == eos_idx] = sep_idx

    embedding, _ = bert_model.bert(bert_labels, output_all_encoded_layers=True)
    # sum over all layers embedding
    embedding = torch.stack(embedding).sum(0)
    # get rid of cls
    embedding = embedding[:, 1:]

    assert labels.shape == embedding.shape[:-1]

    return embedding

def
generate_embedding
(
bert_model
,
labels
)
:
batch_size
,
time
=
labels
.
shape
cls_ids
=
torch
.
full
(
(
batch_size
,
1
)
,
bert_model
.
bert_text_encoder
.
cls_idx
,
dtype
=
labels
.
dtype
,
device
=
labels
.
device
)
bert_labels
=
torch
.
cat
(
[
cls_ids
,
labels
]
,
1
)
# replace eos with sep
eos_idx
=
bert_model
.
bert_text_encoder
.
eos_idx
sep_idx
=
bert_model
.
bert_text_encoder
.
sep_idx
bert_labels
[
bert_labels
==
eos_idx
]
=
sep_idx
embedding
,
_
=
bert_model
.
bert
(
bert_labels
,
output_all_encoded_layers
=
True
)
# sum over all layers embedding
embedding
=
torch
.
stack
(
embedding
)
.
sum
(
0
)
# get rid of cls
embedding
=
embedding
[
:
,
1
:
]
assert
labels
.
shape
==
embedding
.
shape
[
:
1
]
return
embedding


Load fine - tuned bert model given text encoder and checkpoint path .

def load_fine_tuned_model(bert_model, text_encoder, path):
    """Load fine-tuned bert model given text encoder and checkpoint path."""
    bert_text_encoder = BertLikeSentencePieceTextEncoder(text_encoder)

    model = BertForMaskedLM.from_pretrained(bert_model)
    model.bert_text_encoder = bert_text_encoder
    model.bert.embeddings.word_embeddings = nn.Embedding(
        bert_text_encoder.vocab_size, model.bert.embeddings.word_embeddings.weight.shape[1])
    model.config.vocab_size = bert_text_encoder.vocab_size
    model.cls = BertOnlyMLMHead(
        model.config, model.bert.embeddings.word_embeddings.weight)

    model.load_state_dict(torch.load(path))

    return model

def load_fine_tuned_model ( bert_model , text_encoder , path ) : bert_text_encoder = BertLikeSentencePieceTextEncoder ( text_encoder ) model = BertForMaskedLM . from_pretrained ( bert_model ) model . bert_text_encoder = bert_text_encoder model . bert . embeddings . word_embeddings = nn . Embedding ( bert_text_encoder . vocab_size , model . bert . embeddings . word_embeddings . weight . shape [ 1 ] ) model . config . vocab_size = bert_text_encoder . vocab_size model . cls = BertOnlyMLMHead ( model . config , model . bert . embeddings . word_embeddings . weight ) model . load_state_dict ( torch . load ( path ) ) return model


Standard backward step with self . timer and debugger Arguments loss - the loss to perform loss . backward ()

def backward(self, loss):
        '''
        Standard backward step with self.timer and debugger
        Arguments
            loss - the loss to perform loss.backward()
        '''
        self.timer.set()
        loss.backward()
        grad_norm = torch.nn.utils.clip_grad_norm_(
            self.model.parameters(), self.GRAD_CLIP)
        if math.isnan(grad_norm):
            self.verbose('Error : grad norm is NaN @ step '+str(self.step))
        else:
            self.optimizer.step()
        self.timer.cnt('bw')
        return grad_norm

def
backward
(
self
,
loss
)
:
self
.
timer
.
set
(
)
loss
.
backward
(
)
grad_norm
=
torch
.
nn
.
utils
.
clip_grad_norm_
(
self
.
model
.
parameters
(
)
,
self
.
GRAD_CLIP
)
if
math
.
isnan
(
grad_norm
)
:
self
.
verbose
(
'Error : grad norm is NaN @ step '
+
str
(
self
.
step
)
)
else
:
self
.
optimizer
.
step
(
)
self
.
timer
.
cnt
(
'bw'
)
return
grad_norm


Load ckpt if -- load option is specified

def load_ckpt(self):
        ''' Load ckpt if --load option is specified '''
        if self.paras.load:
            # Load weights
            ckpt = torch.load(
                self.paras.load, map_location=self.device if self.mode == 'train' else 'cpu')
            self.model.load_state_dict(ckpt['model'])
            if self.emb_decoder is not None:
                self.emb_decoder.load_state_dict(ckpt['emb_decoder'])
            # if self.amp:
            #    amp.load_state_dict(ckpt['amp'])
            # Load task-dependent items
            metric = "None"
            score = 0.0
            for k, v in ckpt.items():
                if type(v) is float:
                    metric, score = k, v
            if self.mode == 'train':
                self.step = ckpt['global_step']
                self.optimizer.load_opt_state_dict(ckpt['optimizer'])
                self.verbose('Load ckpt from {}, restarting at step {} (recorded {} = {:.2f} %)'.format(
                              self.paras.load, self.step, metric, score))
            else:
                self.model.eval()
                if self.emb_decoder is not None:
                    self.emb_decoder.eval()
                self.verbose('Evaluation target = {} (recorded {} = {:.2f} %)'.format(self.paras.load, metric, score))


def verbose(self, msg):
        ''' Verbose function for print information to stdout'''
        if self.paras.verbose:
            if type(msg) == list:
                for m in msg:
                    print('[INFO]', m.ljust(100))
            else:
                print('[INFO]', msg.ljust(100))

def verbose ( self , msg ) : if self . paras . verbose : if type ( msg ) == list : for m in msg : print ( '[INFO]' , m . ljust ( 100 ) ) else : print ( '[INFO]' , msg . ljust ( 100 ) )


Verbose function for updating progress on stdout ( do not include newline )

def progress(self, msg):
        ''' Verbose function for updating progress on stdout (do not include newline) '''
        if self.paras.verbose:
            sys.stdout.write("\033[K")  # Clear line
            print('[{}] {}'.format(human_format(self.step), msg), end='\r')

def
progress
(
self
,
msg
)
:
if
self
.
paras
.
verbose
:
sys
.
stdout
.
write
(
"\033[K"
)
# Clear line
print
(
'[{}] {}'
.
format
(
human_format
(
self
.
step
)
,
msg
)
,
end
=
'\r'
)


Write log to TensorBoard log_name - <str > Name of tensorboard variable log_value - <dict > / <array > Value of variable ( e . g . dict of losses ) passed if value = None

def write_log(self, log_name, log_dict):
        '''
        Write log to TensorBoard
            log_name  - <str> Name of tensorboard variable
            log_value - <dict>/<array> Value of variable (e.g. dict of losses), passed if value = None
        '''
        if type(log_dict) is dict:
            log_dict = {key: val for key, val in log_dict.items() if (
                val is not None and not math.isnan(val))}
        if log_dict is None:
            pass
        elif len(log_dict) > 0:
            if 'align' in log_name or 'spec' in log_name:
                img, form = log_dict
                self.log.add_image(
                    log_name, img, global_step=self.step, dataformats=form)
            elif 'text' in log_name or 'hyp' in log_name:
                self.log.add_text(log_name, log_dict, self.step)
            else:
                self.log.add_scalars(log_name, log_dict, self.step)

def write_log ( self , log_name , log_dict ) : if type ( log_dict ) is dict : log_dict = { key : val for key , val in log_dict . items ( ) if ( val is not None and not math . isnan ( val ) ) } if log_dict is None : pass elif len ( log_dict ) > 0 : if 'align' in log_name or 'spec' in log_name : img , form = log_dict self . log . add_image ( log_name , img , global_step = self . step , dataformats = form ) elif 'text' in log_name or 'hyp' in log_name : self . log . add_text ( log_name , log_dict , self . step ) else : self . log . add_scalars ( log_name , log_dict , self . step )


Ckpt saver f_name - <str > the name phnof ckpt file ( w / o prefix ) to store overwrite if existed score - <float > The value of metric used to evaluate model

def save_checkpoint(self, f_name, metric, score, show_msg=True):
        ''''
        Ckpt saver
            f_name - <str> the name phnof ckpt file (w/o prefix) to store, overwrite if existed
            score  - <float> The value of metric used to evaluate model
        '''
        ckpt_path = os.path.join(self.ckpdir, f_name)
        full_dict = {
            "model": self.model.state_dict(),
            "optimizer": self.optimizer.get_opt_state_dict(),
            "global_step": self.step,
            metric: score
        }
        # Additional modules to save
        # if self.amp:
        #    full_dict['amp'] = self.amp_lib.state_dict()
        if self.emb_decoder is not None:
            full_dict['emb_decoder'] = self.emb_decoder.state_dict()

        torch.save(full_dict, ckpt_path)
        if show_msg:
            self.verbose("Saved checkpoint (step = {}, {} = {:.2f}) and status @ {}".
                         format(human_format(self.step), metric, score, ckpt_path))


Called by main to load all data After this call data related attributes should be setup ( e . g . self . tr_set self . dev_set ) No return value

def load_data(self):
        '''
        Called by main to load all data
        After this call, data related attributes should be setup (e.g. self.tr_set, self.dev_set)
        No return value
        '''
        raise NotImplementedError

def load_data ( self ) : raise NotImplementedError


Called by main to set models After this call model related attributes should be setup ( e . g . self . l2_loss ) The followings MUST be setup - self . model ( torch . nn . Module ) - self . optimizer ( src . Optimizer ) init . w / self . optimizer = src . Optimizer ( self . model . parameters () ** self . config [ hparas ] ) Loading pre - trained model should also be performed here No return value

def set_model(self):
        '''
        Called by main to set models
        After this call, model related attributes should be setup (e.g. self.l2_loss)
        The followings MUST be setup
            - self.model (torch.nn.Module)
            - self.optimizer (src.Optimizer),
                init. w/ self.optimizer = src.Optimizer(self.model.parameters(),**self.config['hparas'])
        Loading pre-trained model should also be performed here
        No return value
        '''
        raise NotImplementedError

def set_model ( self ) : raise NotImplementedError


Called by main to execute training / inference

def exec(self):
        '''
        Called by main to execute training/inference
        '''
        raise NotImplementedError

def exec ( self ) : raise NotImplementedError


def avgScore(self):
        '''Return the averaged log probability of hypothesis'''
        assert len(self.output_scores) != 0
        return sum(self.output_scores) / len(self.output_scores)

def avgScore ( self ) : assert len ( self . output_scores ) != 0 return sum ( self . output_scores ) / len ( self . output_scores )


def addTopk(self, topi, topv, decoder_state, att_map=None,
                lm_state=None, ctc_state=None, ctc_prob=0.0, ctc_candidates=[]):
        '''Expand current hypothesis with a given beam size'''
        new_hypothesis = []
        term_score = None
        ctc_s, ctc_p = None, None
        beam_size = topi.shape[-1]

        for i in range(beam_size):
            # Detect <eos>
            if topi[i].item() == 1:
                term_score = topv[i].cpu()
                continue

            idxes = self.output_seq[:]     # pass by value
            scores = self.output_scores[:]  # pass by value
            idxes.append(topi[i].cpu())
            scores.append(topv[i].cpu())
            if ctc_state is not None:
                # ToDo: Handle out-of-candidate case.
                idx = ctc_candidates.index(topi[i].item())
                ctc_s = ctc_state[idx, :, :]
                ctc_p = ctc_prob[idx]
            new_hypothesis.append(Hypothesis(decoder_state,
                                             output_seq=idxes, output_scores=scores, lm_state=lm_state,
                                             ctc_state=ctc_s, ctc_prob=ctc_p, att_map=att_map))
        if term_score is not None:
            self.output_seq.append(torch.tensor(1))
            self.output_scores.append(term_score)
            return self, new_hypothesis
        return None, new_hypothesis

def
addTopk
(
self
,
topi
,
topv
,
decoder_state
,
att_map
=
,
lm_state
=
,
ctc_state
=
,
ctc_prob
=
0.0
,
ctc_candidates
=
[
]
)
:
new_hypothesis
=
[
]
term_score
=
ctc_s
,
ctc_p
=
,
beam_size
=
topi
.
shape
[
1
]
for
i
in
range
(
beam_size
)
:
# Detect <eos>
if
topi
[
i
]
.
item
(
)
==
1
:
term_score
=
topv
[
i
]
.
cpu
(
)
continue
idxes
=
self
.
output_seq
[
:
]
# pass by value
scores
=
self
.
output_scores
[
:
]
# pass by value
idxes
.
append
(
topi
[
i
]
.
cpu
(
)
)
scores
.
append
(
topv
[
i
]
.
cpu
(
)
)
if
ctc_state
is
not
:
# ToDo: Handle out-of-candidate case.
idx
=
ctc_candidates
.
index
(
topi
[
i
]
.
item
(
)
)
ctc_s
=
ctc_state
[
idx
,
:
,
:
]
ctc_p
=
ctc_prob
[
idx
]
new_hypothesis
.
append
(
Hypothesis
(
decoder_state
,
output_seq
=
idxes
,
output_scores
=
scores
,
lm_state
=
lm_state
,
ctc_state
=
ctc_s
,
ctc_prob
=
ctc_p
,
att_map
=
att_map
)
)
if
term_score
is
not
:
self
.
output_seq
.
append
(
torch
.
tensor
(
1
)
)
self
.
output_scores
.
append
(
term_score
)
return
self
,
new_hypothesis
return
,
new_hypothesis


def fuse_prob(self, x_emb, dec_logit):
        ''' Takes context and decoder logit to perform word embedding fusion '''
        # Compute distribution for dec/emb
        if self.fuse_normalize:
            emb_logit = nn.functional.linear(nn.functional.normalize(x_emb, dim=-1),
                                             nn.functional.normalize(self.emb_table.weight, dim=-1))
        else:
            emb_logit = nn.functional.linear(x_emb, self.emb_table.weight)
        emb_prob = (nn.functional.relu(self.temp)*emb_logit).softmax(dim=-1)
        dec_prob = dec_logit.softmax(dim=-1)
        # Mix distribution
        if self.fuse_learnable:
            fused_prob = (1-torch.sigmoid(self.fuse_lambda))*dec_prob +\
                torch.sigmoid(self.fuse_lambda)*emb_prob
        else:
            fused_prob = (1-self.fuse_lambda)*dec_prob + \
                self.fuse_lambda*emb_prob
        # Log-prob
        log_fused_prob = (fused_prob+self.eps).log()

        return log_fused_prob

def
fuse_prob
(
self
,
x_emb
,
dec_logit
)
:
# Compute distribution for dec/emb
if
self
.
fuse_normalize
:
emb_logit
=
nn
.
functional
.
linear
(
nn
.
functional
.
normalize
(
x_emb
,
dim
=
1
)
,
nn
.
functional
.
normalize
(
self
.
emb_table
.
weight
,
dim
=
1
)
)
else
:
emb_logit
=
nn
.
functional
.
linear
(
x_emb
,
self
.
emb_table
.
weight
)
emb_prob
=
(
nn
.
functional
.
relu
(
self
.
temp
)
*
emb_logit
)
.
softmax
(
dim
=
1
)
dec_prob
=
dec_logit
.
softmax
(
dim
=
1
)
# Mix distribution
if
self
.
fuse_learnable
:
fused_prob
=
(
1
torch
.
sigmoid
(
self
.
fuse_lambda
)
)
*
dec_prob
+
torch
.
sigmoid
(
self
.
fuse_lambda
)
*
emb_prob
else
:
fused_prob
=
(
1
self
.
fuse_lambda
)
*
dec_prob
+
self
.
fuse_lambda
*
emb_prob
# Log-prob
log_fused_prob
=
(
fused_prob
+
self
.
eps
)
.
log
(
)
return
log_fused_prob


Collects a batch should be list of tuples ( audio_path <str > list of int token <list > ) e . g . [ ( file1 txt1 ) ( file2 txt2 ) ... ]

def collect_audio_batch(batch, audio_transform, mode):
    '''Collects a batch, should be list of tuples (audio_path <str>, list of int token <list>)
       e.g. [(file1,txt1),(file2,txt2),...] '''

    # Bucketed batch should be [[(file1,txt1),(file2,txt2),...]]
    if type(batch[0]) is not tuple:
        batch = batch[0]
    # Make sure that batch size is reasonable
    first_len = audio_transform(str(batch[0][0])).shape[0]
    if first_len > HALF_BATCHSIZE_AUDIO_LEN and mode == 'train':
        batch = batch[:len(batch)//2]

    # Read batch
    file, audio_feat, audio_len, text = [], [], [], []
    with torch.no_grad():
        for b in batch:
            file.append(str(b[0]).split('/')[-1].split('.')[0])
            feat = audio_transform(str(b[0]))
            audio_feat.append(feat)
            audio_len.append(len(feat))
            text.append(torch.LongTensor(b[1]))
    # Descending audio length within each batch
    audio_len, file, audio_feat, text = zip(*[(feat_len, f_name, feat, txt)
                                              for feat_len, f_name, feat, txt in sorted(zip(audio_len, file, audio_feat, text), reverse=True, key=lambda x:x[0])])
    # Zero-padding
    audio_feat = pad_sequence(audio_feat, batch_first=True)
    text = pad_sequence(text, batch_first=True)
    audio_len = torch.LongTensor(audio_len)

    return file, audio_feat, audio_len, text


Collects a batch of text should be list of list of int token e . g . [ txt1 <list > txt2 <list > ... ]

def collect_text_batch(batch, mode):
    '''Collects a batch of text, should be list of list of int token
       e.g. [txt1 <list>,txt2 <list>,...] '''

    # Bucketed batch should be [[txt1, txt2,...]]
    if type(batch[0][0]) is list:
        batch = batch[0]
    # Half batch size if input to long
    if len(batch[0]) > HALF_BATCHSIZE_TEXT_LEN and mode == 'train':
        batch = batch[:len(batch)//2]
    # Read batch
    text = [torch.LongTensor(b) for b in batch]
    # Zero-padding
    text = pad_sequence(text, batch_first=True)

    return text


def create_dataset(tokenizer, ascending, name, path, bucketing, batch_size,
                   train_split=None, dev_split=None, test_split=None):
    ''' Interface for creating all kinds of dataset'''

    # Recognize corpus
    if name.lower() == "librispeech":
        from corpus.librispeech import LibriDataset as Dataset
    else:
        raise NotImplementedError

    # Create dataset
    if train_split is not None:
        # Training mode
        mode = 'train'
        tr_loader_bs = 1 if bucketing and (not ascending) else batch_size
        bucket_size = batch_size if bucketing and (
            not ascending) else 1  # Ascending without bucketing
        # Do not use bucketing for dev set
        dv_set = Dataset(path, dev_split, tokenizer, 1)
        tr_set = Dataset(path, train_split, tokenizer,
                         bucket_size, ascending=ascending)
        # Messages to show
        msg_list = _data_msg(name, path, train_split.__str__(), len(tr_set),
                             dev_split.__str__(), len(dv_set), batch_size, bucketing)

        return tr_set, dv_set, tr_loader_bs, batch_size, mode, msg_list
    else:
        # Testing model
        mode = 'test'
        # Do not use bucketing for dev set
        dv_set = Dataset(path, dev_split, tokenizer, 1)
        # Do not use bucketing for test set
        tt_set = Dataset(path, test_split, tokenizer, 1)
        # Messages to show
        msg_list = _data_msg(name, path, dev_split.__str__(), len(dv_set),
                             test_split.__str__(), len(tt_set), batch_size, False)
        msg_list = [m.replace('Dev', 'Test').replace(
            'Train', 'Dev') for m in msg_list]
        return dv_set, tt_set, batch_size, batch_size, mode, msg_list

def
create_dataset
(
tokenizer
,
ascending
,
name
,
path
,
bucketing
,
batch_size
,
train_split
=
,
dev_split
=
,
test_split
=
)
:
# Recognize corpus
if
name
.
lower
(
)
==
"librispeech"
:
from
corpus
.
librispeech
import
LibriDataset
as
Dataset
else
:
raise
NotImplementedError
# Create dataset
if
train_split
is
not
:
# Training mode
mode
=
'train'
tr_loader_bs
=
1
if
bucketing
and
(
not
ascending
)
else
batch_size
bucket_size
=
batch_size
if
bucketing
and
(
not
ascending
)
else
1
# Ascending without bucketing
# Do not use bucketing for dev set
dv_set
=
Dataset
(
path
,
dev_split
,
tokenizer
,
1
)
tr_set
=
Dataset
(
path
,
train_split
,
tokenizer
,
bucket_size
,
ascending
=
ascending
)
# Messages to show
msg_list
=
_data_msg
(
name
,
path
,
train_split
.
__str__
(
)
,
len
(
tr_set
)
,
dev_split
.
__str__
(
)
,
len
(
dv_set
)
,
batch_size
,
bucketing
)
return
tr_set
,
dv_set
,
tr_loader_bs
,
batch_size
,
mode
,
msg_list
else
:
# Testing model
mode
=
'test'
# Do not use bucketing for dev set
dv_set
=
Dataset
(
path
,
dev_split
,
tokenizer
,
1
)
# Do not use bucketing for test set
tt_set
=
Dataset
(
path
,
test_split
,
tokenizer
,
1
)
# Messages to show
msg_list
=
_data_msg
(
name
,
path
,
dev_split
.
__str__
(
)
,
len
(
dv_set
)
,
test_split
.
__str__
(
)
,
len
(
tt_set
)
,
batch_size
,
False
)
msg_list
=
[
m
.
replace
(
'Dev'
,
'Test'
)
.
replace
(
'Train'
,
'Dev'
)
for
m
in
msg_list
]
return
dv_set
,
tt_set
,
batch_size
,
batch_size
,
mode
,
msg_list


def create_textset(tokenizer, train_split, dev_split, name, path, bucketing, batch_size):
    ''' Interface for creating all kinds of text dataset'''
    msg_list = []

    # Recognize corpus
    if name.lower() == "librispeech":
        from corpus.librispeech import LibriTextDataset as Dataset
    else:
        raise NotImplementedError

    # Create dataset
    bucket_size = batch_size if bucketing else 1
    tr_loader_bs = 1 if bucketing else batch_size
    # Do not use bucketing for dev set
    dv_set = Dataset(path, dev_split, tokenizer, 1)
    tr_set = Dataset(path, train_split, tokenizer, bucket_size)

    # Messages to show
    msg_list = _data_msg(name, path, train_split.__str__(), len(tr_set),
                         dev_split.__str__(), len(dv_set), batch_size, bucketing)

    return tr_set, dv_set, tr_loader_bs, batch_size, msg_list

def
create_textset
(
tokenizer
,
train_split
,
dev_split
,
name
,
path
,
bucketing
,
batch_size
)
:
msg_list
=
[
]
# Recognize corpus
if
name
.
lower
(
)
==
"librispeech"
:
from
corpus
.
librispeech
import
LibriTextDataset
as
Dataset
else
:
raise
NotImplementedError
# Create dataset
bucket_size
=
batch_size
if
bucketing
else
1
tr_loader_bs
=
1
if
bucketing
else
batch_size
# Do not use bucketing for dev set
dv_set
=
Dataset
(
path
,
dev_split
,
tokenizer
,
1
)
tr_set
=
Dataset
(
path
,
train_split
,
tokenizer
,
bucket_size
)
# Messages to show
msg_list
=
_data_msg
(
name
,
path
,
train_split
.
__str__
(
)
,
len
(
tr_set
)
,
dev_split
.
__str__
(
)
,
len
(
dv_set
)
,
batch_size
,
bucketing
)
return
tr_set
,
dv_set
,
tr_loader_bs
,
batch_size
,
msg_list


Prepare dataloader for training / validation

def load_dataset(n_jobs, use_gpu, pin_memory, ascending, corpus, audio, text):
    ''' Prepare dataloader for training/validation'''

    # Audio feature extractor
    audio_transform, feat_dim = create_transform(audio.copy())
    # Text tokenizer
    tokenizer = load_text_encoder(**text)
    # Dataset (in testing mode, tr_set=dv_set, dv_set=tt_set)
    tr_set, dv_set, tr_loader_bs, dv_loader_bs, mode, data_msg = create_dataset(
        tokenizer, ascending, **corpus)
    # Collect function
    collect_tr = partial(collect_audio_batch,
                         audio_transform=audio_transform, mode=mode)
    collect_dv = partial(collect_audio_batch,
                         audio_transform=audio_transform, mode='test')
    # Shuffle/drop applied to training set only
    shuffle = (mode == 'train' and not ascending)
    drop_last = shuffle
    # Create data loader
    tr_set = DataLoader(tr_set, batch_size=tr_loader_bs, shuffle=shuffle, drop_last=drop_last, collate_fn=collect_tr,
                        num_workers=n_jobs, pin_memory=use_gpu)
    dv_set = DataLoader(dv_set, batch_size=dv_loader_bs, shuffle=False, drop_last=False, collate_fn=collect_dv,
                        num_workers=n_jobs, pin_memory=pin_memory)
    # Messages to show
    data_msg.append('I/O spec.  | Audio feature = {}\t| feature dim = {}\t| Token type = {}\t| Vocab size = {}'
                    .format(audio['feat_type'], feat_dim, tokenizer.token_type, tokenizer.vocab_size))

    return tr_set, dv_set, feat_dim, tokenizer.vocab_size, tokenizer, data_msg


def _data_msg(name, path, train_split, tr_set, dev_split, dv_set, batch_size, bucketing):
    ''' List msg for verbose function '''
    msg_list = []
    msg_list.append('Data spec. | Corpus = {} (from {})'.format(name, path))
    msg_list.append('           | Train sets = {}\t| Number of utts = {}'.format(
        train_split, tr_set))
    msg_list.append(
        '           | Dev sets = {}\t| Number of utts = {}'.format(dev_split, dv_set))
    msg_list.append('           | Batch size = {}\t\t| Bucketing = {}'.format(
        batch_size, bucketing))
    return msg_list


def set_state(self, prev_state, prev_attn):
        ''' Setting up all memory states for beam decoding'''
        self.decoder.set_state(prev_state)
        self.attention.set_mem(prev_attn)

def set_state ( self , prev_state , prev_attn ) : self . decoder . set_state ( prev_state ) self . attention . set_mem ( prev_attn )


Arguments audio_feature - [ BxTxD ] Acoustic feature with shape feature_len - [ B ] Length of each sample in a batch decode_step - [ int ] The maximum number of attention decoder steps tf_rate - [ 0 1 ] The probability to perform teacher forcing for each step teacher - [ BxL ] Ground truth for teacher forcing with sentence length L emb_decoder - [ obj ] Introduces the word embedding decoder different behavior for training / inference At training stage this ONLY affects self - sampling ( output remains the same ) At inference stage this affects output to become log prob . with distribution fusion get_dec_state - [ bool ] If true return decoder state [ BxLxD ] for other purpose

def forward(self, audio_feature, feature_len, decode_step, tf_rate=0.0, teacher=None,
                emb_decoder=None, get_dec_state=False):
        '''
        Arguments
            audio_feature - [BxTxD] Acoustic feature with shape
            feature_len   - [B]     Length of each sample in a batch
            decode_step   - [int]   The maximum number of attention decoder steps
            tf_rate       - [0,1]   The probability to perform teacher forcing for each step
            teacher       - [BxL] Ground truth for teacher forcing with sentence length L
            emb_decoder   - [obj]   Introduces the word embedding decoder, different behavior for training/inference
                                    At training stage, this ONLY affects self-sampling (output remains the same)
                                    At inference stage, this affects output to become log prob. with distribution fusion
            get_dec_state - [bool]  If true, return decoder state [BxLxD] for other purpose
        '''
        # Init
        bs = audio_feature.shape[0]
        ctc_output, att_output, att_seq = None, None, None
        dec_state = [] if get_dec_state else None

        # Encode
        encode_feature, encode_len = self.encoder(audio_feature, feature_len)

        # CTC based decoding
        if self.enable_ctc:
            ctc_output = F.log_softmax(self.ctc_layer(encode_feature), dim=-1)

        # Attention based decoding
        if self.enable_att:
            # Init (init char = <SOS>, reset all rnn state and cell)
            self.decoder.init_state(bs)
            self.attention.reset_mem()
            last_char = self.pre_embed(torch.zeros(
                (bs), dtype=torch.long, device=encode_feature.device))
            att_seq, output_seq = [], []

            # Preprocess data for teacher forcing
            if teacher is not None:
                teacher = self.embed_drop(self.pre_embed(teacher))

            # Decode
            for t in range(decode_step):
                # Attend (inputs current state of first layer, encoded features)
                attn, context = self.attention(
                    self.decoder.get_query(), encode_feature, encode_len)
                # Decode (inputs context + embedded last character)
                decoder_input = torch.cat([last_char, context], dim=-1)
                cur_char, d_state = self.decoder(decoder_input)
                # Prepare output as input of next step
                if (teacher is not None):
                    # Training stage
                    if (tf_rate == 1) or (torch.rand(1).item() <= tf_rate):
                        # teacher forcing
                        last_char = teacher[:, t, :]
                    else:
                        # self-sampling (replace by argmax may be another choice)
                        with torch.no_grad():
                            if (emb_decoder is not None) and emb_decoder.apply_fuse:
                                _, cur_prob = emb_decoder(
                                    d_state, cur_char, return_loss=False)
                            else:
                                cur_prob = cur_char.softmax(dim=-1)
                            sampled_char = Categorical(cur_prob).sample()
                        last_char = self.embed_drop(
                            self.pre_embed(sampled_char))
                else:
                    # Inference stage
                    if (emb_decoder is not None) and emb_decoder.apply_fuse:
                        _, cur_char = emb_decoder(
                            d_state, cur_char, return_loss=False)
                    # argmax for inference
                    last_char = self.pre_embed(torch.argmax(cur_char, dim=-1))

                # save output of each step
                output_seq.append(cur_char)
                att_seq.append(attn)
                if get_dec_state:
                    dec_state.append(d_state)

            att_output = torch.stack(output_seq, dim=1)  # BxTxV
            att_seq = torch.stack(att_seq, dim=2)       # BxNxDtxT
            if get_dec_state:
                dec_state = torch.stack(dec_state, dim=1)

        return ctc_output, encode_len, att_output, att_seq, dec_state


def init_state(self, bs):
        ''' Set all hidden states to zeros '''
        device = next(self.parameters()).device
        if self.enable_cell:
            self.hidden_state = (torch.zeros((self.layer, bs, self.dim), device=device),
                                 torch.zeros((self.layer, bs, self.dim), device=device))
        else:
            self.hidden_state = torch.zeros(
                (self.layer, bs, self.dim), device=device)
        return self.get_state()

def init_state ( self , bs ) : device = next ( self . parameters ( ) ) . device if self . enable_cell : self . hidden_state = ( torch . zeros ( ( self . layer , bs , self . dim ) , device = device ) , torch . zeros ( ( self . layer , bs , self . dim ) , device = device ) ) else : self . hidden_state = torch . zeros ( ( self . layer , bs , self . dim ) , device = device ) return self . get_state ( )


Set all hidden states / cells for decoding purpose

def set_state(self, hidden_state):
        ''' Set all hidden states/cells, for decoding purpose'''
        device = next(self.parameters()).device
        if self.enable_cell:
            self.hidden_state = (hidden_state[0].to(
                device), hidden_state[1].to(device))
        else:
            self.hidden_state = hidden_state.to(device)

def set_state ( self , hidden_state ) : device = next ( self . parameters ( ) ) . device if self . enable_cell : self . hidden_state = ( hidden_state [ 0 ] . to ( device ) , hidden_state [ 1 ] . to ( device ) ) else : self . hidden_state = hidden_state . to ( device )


Return all hidden states / cells for decoding purpose

def get_state(self):
        ''' Return all hidden states/cells, for decoding purpose'''
        if self.enable_cell:
            return (self.hidden_state[0].cpu(), self.hidden_state[1].cpu())
        else:
            return self.hidden_state.cpu()

def get_state ( self ) : if self . enable_cell : return ( self . hidden_state [ 0 ] . cpu ( ) , self . hidden_state [ 1 ] . cpu ( ) ) else : return self . hidden_state . cpu ( )


def get_query(self):
        ''' Return state of all layers as query for attention '''
        if self.enable_cell:
            return self.hidden_state[0].transpose(0, 1).reshape(-1, self.dim*self.layer)
        else:
            return self.hidden_state.transpose(0, 1).reshape(-1, self.dim*self.layer)

def get_query ( self ) : if self . enable_cell : return self . hidden_state [ 0 ] . transpose ( 0 , 1 ) . reshape ( - 1 , self . dim * self . layer ) else : return self . hidden_state . transpose ( 0 , 1 ) . reshape ( - 1 , self . dim * self . layer )


def forward(self, x):
        ''' Decode and transform into vocab '''
        if not self.training:
            self.layers.flatten_parameters()
        x, self.hidden_state = self.layers(x.unsqueeze(1), self.hidden_state)
        x = x.squeeze(1)
        char = self.char_trans(self.final_dropout(x))
        return char, x

def forward ( self , x ) : if not self . training : self . layers . flatten_parameters ( ) x , self . hidden_state = self . layers ( x . unsqueeze ( 1 ) , self . hidden_state ) x = x . squeeze ( 1 ) char = self . char_trans ( self . final_dropout ( x ) ) return char , x


Given prefix g return the probability of all possible sequence y ( where y = concat ( g c )) This function computes all possible tokens for c ( memory inefficient )

def full_compute(self, g, r_prev):
        '''Given prefix g, return the probability of all possible sequence y (where y = concat(g,c))
           This function computes all possible tokens for c (memory inefficient)'''
        prefix_length = len(g)
        last_char = g[-1] if prefix_length > 0 else 0

        # init. r
        r = np.full((self.input_length, 2, self.odim),
                    self.logzero, dtype=np.float32)

        # start from len(g) because is impossible for CTC to generate |y|>|X|
        start = max(1, prefix_length)

        if prefix_length == 0:
            r[0, 0, :] = self.x[0, :]    # if g = <sos>

        psi = r[start-1, 0, :]

        phi = np.logaddexp(r_prev[:, 0], r_prev[:, 1])

        for t in range(start, self.input_length):
            # prev_blank
            prev_blank = np.full((self.odim), r_prev[t-1, 1], dtype=np.float32)
            # prev_nonblank
            prev_nonblank = np.full(
                (self.odim), r_prev[t-1, 0], dtype=np.float32)
            prev_nonblank[last_char] = self.logzero

            phi = np.logaddexp(prev_nonblank, prev_blank)
            # P(h|current step is non-blank) = [ P(prev. step = y) + P()]*P(c)
            r[t, 0, :] = np.logaddexp(r[t-1, 0, :], phi) + self.x[t, :]
            # P(h|current step is blank) = [P(prev. step is blank) + P(prev. step is non-blank)]*P(now=blank)
            r[t, 1, :] = np.logaddexp(
                r[t-1, 1, :], r[t-1, 0, :]) + self.x[t, self.blank]
            psi = np.logaddexp(psi, phi+self.x[t, :])

        #psi[self.eos] = np.logaddexp(r_prev[-1,0], r_prev[-1,1])
        return psi, np.rollaxis(r, 2)


Given prefix g return the probability of all possible sequence y ( where y = concat ( g c )) This function considers only those tokens in candidates for c ( memory efficient )

def cheap_compute(self, g, r_prev, candidates):
        '''Given prefix g, return the probability of all possible sequence y (where y = concat(g,c))
           This function considers only those tokens in candidates for c (memory efficient)'''
        prefix_length = len(g)
        odim = len(candidates)
        last_char = g[-1] if prefix_length > 0 else 0

        # init. r
        r = np.full((self.input_length, 2, len(candidates)),
                    self.logzero, dtype=np.float32)

        # start from len(g) because is impossible for CTC to generate |y|>|X|
        start = max(1, prefix_length)

        if prefix_length == 0:
            r[0, 0, :] = self.x[0, candidates]    # if g = <sos>

        psi = r[start-1, 0, :]
        # Phi = (prev_nonblank,prev_blank)
        sum_prev = np.logaddexp(r_prev[:, 0], r_prev[:, 1])
        phi = np.repeat(sum_prev[..., None],odim,axis=-1)
        # Handle edge case : last tok of prefix in candidates
        if  prefix_length>0 and last_char in candidates:
            phi[:,candidates.index(last_char)] = r_prev[:,1]

        for t in range(start, self.input_length):
            # prev_blank
            # prev_blank = np.full((odim), r_prev[t-1, 1], dtype=np.float32)
            # prev_nonblank
            # prev_nonblank = np.full((odim), r_prev[t-1, 0], dtype=np.float32)
            # phi = np.logaddexp(prev_nonblank, prev_blank)
            # P(h|current step is non-blank) =  P(prev. step = y)*P(c)
            r[t, 0, :] = np.logaddexp( r[t-1, 0, :], phi[t-1]) + self.x[t, candidates]
            # P(h|current step is blank) = [P(prev. step is blank) + P(prev. step is non-blank)]*P(now=blank)
            r[t, 1, :] = np.logaddexp( r[t-1, 1, :], r[t-1, 0, :]) + self.x[t, self.blank]
            psi = np.logaddexp(psi, phi[t-1,]+self.x[t, candidates])

        # P(end of sentence) = P(g)
        if self.eos in candidates:
            psi[candidates.index(self.eos)] = sum_prev[-1]
        return psi, np.rollaxis(r, 2)