Update seq2seq.py

Author: iml1111

src/11_seq2seq/modules/seq2seq.py

@@ -3,8 +3,7 @@
 from torch.nn.utils.rnn import pack_padded_sequence as pack
 from torch.nn.utils.rnn import pad_packed_sequence as unpack
 
-import simple_nmt.data_loader as data_loader
-from simple_nmt.search import SingleBeamSearchBoard
+import modules.data_loader as data_loader
 
 
 class Attention(nn.Module):
@@ -370,150 +369,3 @@ def search(self, src, is_greedy=True, max_length=255):
         # |indice| = (batch_size, length)
 
         return y_hats, indice
-
-    #@profile
-    def batch_beam_search(
-        self,
-        src,
-        beam_size=5,
-        max_length=255,
-        n_best=1,
-        length_penalty=.2
-    ):
-        mask, x_length = None, None
-
-        if isinstance(src, tuple):
-            x, x_length = src
-            mask = self.generate_mask(x, x_length)
-            # |mask| = (batch_size, length)
-        else:
-            x = src
-        batch_size = x.size(0)
-
-        emb_src = self.emb_src(x)
-        h_src, h_0_tgt = self.encoder((emb_src, x_length))
-        # |h_src| = (batch_size, length, hidden_size)
-        h_0_tgt = self.fast_merge_encoder_hiddens(h_0_tgt)
-
-        # initialize 'SingleBeamSearchBoard' as many as batch_size
-        boards = [SingleBeamSearchBoard(
-            h_src.device,
-            {
-                'hidden_state': {
-                    'init_status': h_0_tgt[0][:, i, :].unsqueeze(1),
-                    'batch_dim_index': 1,
-                }, # |hidden_state| = (n_layers, batch_size, hidden_size)
-                'cell_state': {
-                    'init_status': h_0_tgt[1][:, i, :].unsqueeze(1),
-                    'batch_dim_index': 1,
-                }, # |cell_state| = (n_layers, batch_size, hidden_size)
-                'h_t_1_tilde': {
-                    'init_status': None,
-                    'batch_dim_index': 0,
-                }, # |h_t_1_tilde| = (batch_size, 1, hidden_size)
-            },
-            beam_size=beam_size,
-            max_length=max_length,
-        ) for i in range(batch_size)]
-        is_done = [board.is_done() for board in boards]
-
-        length = 0
-        # Run loop while sum of 'is_done' is smaller than batch_size, 
-        # or length is still smaller than max_length.
-        while sum(is_done) < batch_size and length <= max_length:
-            # current_batch_size = sum(is_done) * beam_size
-
-            # Initialize fabricated variables.
-            # As far as batch-beam-search is running, 
-            # temporary batch-size for fabricated mini-batch is 
-            # 'beam_size'-times bigger than original batch_size.
-            fab_input, fab_hidden, fab_cell, fab_h_t_tilde = [], [], [], []
-            fab_h_src, fab_mask = [], []
-            
-            # Build fabricated mini-batch in non-parallel way.
-            # This may cause a bottle-neck.
-            for i, board in enumerate(boards):
-                # Batchify if the inference for the sample is still not finished.
-                if board.is_done() == 0:
-                    y_hat_i, prev_status = board.get_batch()
-                    hidden_i    = prev_status['hidden_state']
-                    cell_i      = prev_status['cell_state']
-                    h_t_tilde_i = prev_status['h_t_1_tilde']
-
-                    fab_input  += [y_hat_i]
-                    fab_hidden += [hidden_i]
-                    fab_cell   += [cell_i]
-                    fab_h_src  += [h_src[i, :, :]] * beam_size
-                    fab_mask   += [mask[i, :]] * beam_size
-                    if h_t_tilde_i is not None:
-                        fab_h_t_tilde += [h_t_tilde_i]
-                    else:
-                        fab_h_t_tilde = None
-
-            # Now, concatenate list of tensors.
-            fab_input  = torch.cat(fab_input,  dim=0)
-            fab_hidden = torch.cat(fab_hidden, dim=1)
-            fab_cell   = torch.cat(fab_cell,   dim=1)
-            fab_h_src  = torch.stack(fab_h_src)
-            fab_mask   = torch.stack(fab_mask)
-            if fab_h_t_tilde is not None:
-                fab_h_t_tilde = torch.cat(fab_h_t_tilde, dim=0)
-            # |fab_input|     = (current_batch_size, 1)
-            # |fab_hidden|    = (n_layers, current_batch_size, hidden_size)
-            # |fab_cell|      = (n_layers, current_batch_size, hidden_size)
-            # |fab_h_src|     = (current_batch_size, length, hidden_size)
-            # |fab_mask|      = (current_batch_size, length)
-            # |fab_h_t_tilde| = (current_batch_size, 1, hidden_size)
-
-            emb_t = self.emb_dec(fab_input)
-            # |emb_t| = (current_batch_size, 1, word_vec_size)
-
-            fab_decoder_output, (fab_hidden, fab_cell) = self.decoder(emb_t,
-                                                                      fab_h_t_tilde,
-                                                                      (fab_hidden, fab_cell))
-            # |fab_decoder_output| = (current_batch_size, 1, hidden_size)
-            context_vector = self.attn(fab_h_src, fab_decoder_output, fab_mask)
-            # |context_vector| = (current_batch_size, 1, hidden_size)
-            fab_h_t_tilde = self.tanh(self.concat(torch.cat([fab_decoder_output,
-                                                             context_vector
-                                                             ], dim=-1)))
-            # |fab_h_t_tilde| = (current_batch_size, 1, hidden_size)
-            y_hat = self.generator(fab_h_t_tilde)
-            # |y_hat| = (current_batch_size, 1, output_size)
-
-            # separate the result for each sample.
-            # fab_hidden[:, begin:end, :] = (n_layers, beam_size, hidden_size)
-            # fab_cell[:, begin:end, :]   = (n_layers, beam_size, hidden_size)
-            # fab_h_t_tilde[begin:end]    = (beam_size, 1, hidden_size)
-            cnt = 0
-            for board in boards:
-                if board.is_done() == 0:
-                    # Decide a range of each sample.
-                    begin = cnt * beam_size
-                    end = begin + beam_size
-
-                    # pick k-best results for each sample.
-                    board.collect_result(
-                        y_hat[begin:end],
-                        {
-                            'hidden_state': fab_hidden[:, begin:end, :],
-                            'cell_state'  : fab_cell[:, begin:end, :],
-                            'h_t_1_tilde' : fab_h_t_tilde[begin:end],
-                        },
-                    )
-                    cnt += 1
-
-            is_done = [board.is_done() for board in boards]
-            length += 1
-
-        # pick n-best hypothesis.
-        batch_sentences, batch_probs = [], []
-
-        # Collect the results.
-        for i, board in enumerate(boards):
-            sentences, probs = board.get_n_best(n_best, length_penalty=length_penalty)
-
-            batch_sentences += [sentences]
-            batch_probs     += [probs]
-
-        return batch_sentences, batch_probs