remove completed sentences from batch

            dataset.src, dataset.dst, max_tokens, max_sentences,

            max_positions=max_positions,

            ignore_invalid_inputs=skip_invalid_size_inputs_valid_test,

            descending=descending)

            descending=descending,

            allow_different_src_lens=True)

        batch_sampler = mask_batches(batch_sampler, shard_id=shard_id, num_shards=num_shards)

        return torch.utils.data.DataLoader(

            dataset, num_workers=num_workers, collate_fn=dataset.collater,

def batches_by_size(src, dst, max_tokens=None, max_sentences=None,

                    max_positions=(1024, 1024), ignore_invalid_inputs=False,

                    descending=False, required_batch_size_multiple=1):

                    descending=False, required_batch_size_multiple=1, allow_different_src_lens=False):

    """Returns batches of indices sorted by size. Sequences with different

    source lengths are not allowed in the same batch."""

    assert isinstance(src, IndexedDataset) and (dst is None or isinstance(dst, IndexedDataset))

        indices = np.flip(indices, 0)

    return list(_make_batches(

        src, dst, indices, max_tokens, max_sentences, max_positions,

        ignore_invalid_inputs, allow_different_src_lens=False,

        ignore_invalid_inputs, allow_different_src_lens=allow_different_src_lens,

        required_batch_size_multiple=required_batch_size_multiple,

    ))

        """

        def apply_reorder_incremental_state(module):

            if module != self and hasattr(module, 'reorder_incremental_state'):

                module.reorder_incremental_state(incremental_state, new_order)

                module.reorder_incremental_state(

                    incremental_state,

                    new_order,

                )

        self.apply(apply_reorder_incremental_state)

    def reorder_encoder_out(self, encoder_out, new_order):

        return encoder_out

    def set_beam_size(self, beam_size):

        """Sets the beam size in the decoder and all children."""

        if getattr(self, '_beam_size', -1) != beam_size:

        self.num_attention_layers = None

        num_embeddings = len(dictionary)

        padding_idx = dictionary.pad()

        self.embed_tokens = Embedding(num_embeddings, embed_dim, padding_idx)

        if embed_dict:

            self.embed_tokens = utils.load_embedding(embed_dict, self.dictionary, self.embed_tokens)

        self.padding_idx = dictionary.pad()

        self.embed_tokens = Embedding(num_embeddings, embed_dim, self.padding_idx)

        self.embed_positions = PositionalEmbedding(

            max_positions,

            embed_dim,

            padding_idx,

            self.padding_idx,

            left_pad=LanguagePairDataset.LEFT_PAD_SOURCE,

        )

        # project to size of convolution

        x = self.fc1(x)

        # used to mask padding in input

        encoder_padding_mask = src_tokens.eq(self.padding_idx).t()  # -> T x B

        if not encoder_padding_mask.any():

            encoder_padding_mask = None

        # B x T x C -> T x B x C

        x = x.transpose(0, 1)

        # temporal convolutions

        for proj, conv in zip(self.projections, self.convolutions):

            residual = x if proj is None else proj(x)

            if encoder_padding_mask is not None:

                x = x.masked_fill(encoder_padding_mask.unsqueeze(-1), 0)

            x = F.dropout(x, p=self.dropout, training=self.training)

            if conv.kernel_size[0] % 2 == 1:

                # padding is implicit in the conv

        # project back to size of embedding

        x = self.fc2(x)

        if encoder_padding_mask is not None:

            encoder_padding_mask = encoder_padding_mask.t()  # -> B x T

            x = x.masked_fill(encoder_padding_mask.unsqueeze(-1), 0)

        # scale gradients (this only affects backward, not forward)

        x = GradMultiply.apply(x, 1.0 / (2.0 * self.num_attention_layers))

        # add output to input embedding for attention

        y = (x + input_embedding) * math.sqrt(0.5)

        return x, y

        return {

            'encoder_out': (x, y),

            'encoder_padding_mask': encoder_padding_mask,  # B x T

        }

    def max_positions(self):

        """Maximum input length supported by the encoder."""

        self.bmm = bmm if bmm is not None else torch.bmm

    def forward(self, x, target_embedding, encoder_out):

    def forward(self, x, target_embedding, encoder_out, encoder_padding_mask):

        residual = x

        # attention

        x = (self.in_projection(x) + target_embedding) * math.sqrt(0.5)

        x = self.bmm(x, encoder_out[0])

        # don't attend over padding

        if encoder_padding_mask is not None:

            x = x.float().masked_fill(

                encoder_padding_mask.unsqueeze(1),

                float('-inf')

            ).type_as(x)  # FP16 support: cast to float and back

        # softmax over last dim

        sz = x.size()

        x = F.softmax(x.view(sz[0] * sz[1], sz[2]), dim=1)

        x = self.bmm(x, encoder_out[1])

        # scale attention output

        # scale attention output (respecting potentially different lengths)

        s = encoder_out[1].size(1)

        if encoder_padding_mask is None:

            x = x * (s * math.sqrt(1.0 / s))

        else:

            s = s - encoder_padding_mask.type_as(x).sum(dim=1, keepdim=True)  # exclude padding

            s = s.unsqueeze(-1)

            x = x * (s * s.rsqrt())

        # project back

        x = (self.out_projection(x) + residual) * math.sqrt(0.5)

        else:

            self.fc3 = Linear(out_embed_dim, num_embeddings, dropout=dropout)

    def forward(self, prev_output_tokens, encoder_out, incremental_state=None):

    def forward(self, prev_output_tokens, encoder_out_dict, incremental_state=None):

        encoder_out = encoder_out_dict['encoder_out']

        encoder_padding_mask = encoder_out_dict['encoder_padding_mask']

        # split and transpose encoder outputs

        encoder_a, encoder_b = self._split_encoder_out(encoder_out, incremental_state)

            if attention is not None:

                x = self._transpose_if_training(x, incremental_state)

                x, attn_scores = attention(x, target_embedding, (encoder_a, encoder_b))

                x, attn_scores = attention(x, target_embedding, (encoder_a, encoder_b), encoder_padding_mask)

                attn_scores = attn_scores / num_attn_layers

                if avg_attn_scores is None:

                    avg_attn_scores = attn_scores

        return x

    def reorder_incremental_state(self, incremental_state, new_order):

        super().reorder_incremental_state(incremental_state, new_order)

        encoder_out = utils.get_incremental_state(self, incremental_state, 'encoder_out')

        if encoder_out is not None:

            def update_enc_out(enc_out):

                return enc_out.index_select(0, new_order)

            encoder_out = tuple([update_enc_out(eo) for eo in encoder_out])

            utils.set_incremental_state(self, incremental_state, 'encoder_out', encoder_out)

    def reorder_encoder_out(self, encoder_out_dict, new_order):

        if encoder_out_dict['encoder_padding_mask'] is not None:

            encoder_out_dict['encoder_padding_mask'] = encoder_out_dict['encoder_padding_mask'].index_select(0, new_order)

        return encoder_out_dict

def Embedding(num_embeddings, embedding_dim, padding_idx):

    m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx)

    m.weight.data.normal_(0, 0.1)

                        dim=0).view(bsz, self.output_units))

            return x, bi_final_hiddens, bi_final_cells

        return x, final_hiddens, final_cells

        encoder_padding_mask = src_tokens.eq(self.padding_idx).t()

        return {

            'encoder_out': (x, final_hiddens, final_cells),

            'encoder_padding_mask': encoder_padding_mask if encoder_padding_mask.any() else None

        }

    def max_positions(self):

        """Maximum input length supported by the encoder."""

        self.input_proj = Linear(input_embed_dim, output_embed_dim, bias=False)

        self.output_proj = Linear(2*output_embed_dim, output_embed_dim, bias=False)

    def forward(self, input, source_hids, src_lengths=None):

    def forward(self, input, source_hids, encoder_padding_mask):

        # input: bsz x input_embed_dim

        # source_hids: srclen x bsz x output_embed_dim

        # compute attention

        attn_scores = (source_hids * x.unsqueeze(0)).sum(dim=2)

        # don't attend over padding

        if encoder_padding_mask is not None:

            attn_scores = attn_scores.float().masked_fill_(

                encoder_padding_mask,

                float('-inf')

            ).type_as(attn_scores)  # FP16 support: cast to float and back

        attn_scores = F.softmax(attn_scores.t(), dim=1).t()  # srclen x bsz

        # sum weighted sources

            self.additional_fc = Linear(hidden_size, out_embed_dim)

        self.fc_out = Linear(out_embed_dim, num_embeddings, dropout=dropout_out)

    def forward(self, prev_output_tokens, encoder_out, incremental_state=None):

    def forward(self, prev_output_tokens, encoder_out_dict, incremental_state=None):

        encoder_out = encoder_out_dict['encoder_out']

        encoder_padding_mask = encoder_out_dict['encoder_padding_mask']

        if incremental_state is not None:

            prev_output_tokens = prev_output_tokens[:, -1:]

        bsz, seqlen = prev_output_tokens.size()

            # apply attention using the last layer's hidden state

            if self.attention is not None:

                out, attn_scores[:, j, :] = self.attention(hidden, encoder_outs)

                out, attn_scores[:, j, :] = self.attention(hidden, encoder_outs, encoder_padding_mask)

            else:

                out = hidden

            out = F.dropout(out, p=self.dropout_out, training=self.training)

        new_state = tuple(map(reorder_state, cached_state))

        utils.set_incremental_state(self, incremental_state, 'cached_state', new_state)

    def reorder_encoder_out(self, encoder_out_dict, new_order):

        encoder_out_dict['encoder_out'] = tuple(

            eo.index_select(1, new_order) for eo in encoder_out_dict['encoder_out'])

        if encoder_out_dict['encoder_padding_mask'] is not None:

            encoder_out_dict['encoder_padding_mask'] = encoder_out_dict['encoder_padding_mask'].index_select(1, new_order)

        return encoder_out_dict

    def max_positions(self):

        """Maximum output length supported by the decoder."""

        return int(1e5)  # an arbitrary large number

        return state_dict

    def reorder_encoder_out(self, encoder_out, new_order):

        if encoder_out['encoder_padding_mask'] is not None:

            encoder_out['encoder_padding_mask'] = encoder_out['encoder_padding_mask'].index_select(0, new_order)

        return encoder_out

class TransformerEncoderLayer(nn.Module):

    """Encoder layer block.