Merge pull request #107 from facebookresearch/oss-merge-internal

import torch.nn.functional as F

from . import FairseqCriterion, register_criterion

from fairseq import utils

@register_criterion('cross_entropy')

class CrossEntropyCriterion(FairseqCriterion):

                          reduce=reduce)

        sample_size = sample['target'].size(0) if self.args.sentence_avg else sample['ntokens']

        logging_output = {

            'loss': loss.data[0] if reduce else loss.data,

            'loss': utils.item(loss.data) if reduce else loss.data,

            'ntokens': sample['ntokens'],

            'sample_size': sample_size,

        }

import math

import torch

from torch.autograd import Variable

import torch.nn.functional as F

from fairseq import utils

        norm = grad_input.size(-1)

        if weights is not None:

            if isinstance(grad_input, Variable) and not isinstance(weights, Variable):

                weights = Variable(weights, requires_grad=False)

            norm = weights.sum()

            grad_input.mul(weights.view(1, weights.size(0)).expand_as(grad_input))

    @staticmethod

    def backward(ctx, grad):

        return utils.volatile_variable(ctx.grad_input) * grad, None, None, None, None, None

        grad_input = ctx.grad_input

        if not isinstance(grad_input, torch.autograd.Variable):

            grad_input = utils.volatile_variable(grad_input)

        return grad_input * grad, None, None, None, None, None

@register_criterion('label_smoothed_cross_entropy')

        nll_loss = F.nll_loss(lprobs, target, size_average=False, ignore_index=self.padding_idx, reduce=reduce)

        sample_size = sample['target'].size(0) if self.args.sentence_avg else sample['ntokens']

        logging_output = {

            'loss': loss.data[0] if reduce else loss.data,

            'nll_loss': nll_loss.data[0] if reduce else loss.data,

            'loss': utils.item(loss.data) if reduce else loss.data,

            'nll_loss': utils.item(nll_loss.data) if reduce else loss.data,

            'ntokens': sample['ntokens'],

            'sample_size': sample_size,

        }

    if len(enc) >= max_size:

        raise ValueError('encoded data exceeds max_size: {}'.format(len(enc)))

    in_buffer[0] = len(enc)

    in_buffer[1:len(enc)+1] = torch.ByteTensor(enc)

    in_buffer[1:len(enc)+1] = torch.ByteTensor(list(enc))

    torch.distributed.all_gather(out_buffers, in_buffer.cuda())

    def __init__(self, dictionary):

        super().__init__(dictionary)

        self._is_incremental_eval = False

        self._incremental_state = {}

    def forward(self, prev_output_tokens, encoder_out):

        if self._is_incremental_eval:

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

        raise NotImplementedError

        else:

            raise NotImplementedError

    def incremental_inference(self):

        """Context manager for incremental inference.

        This provides an optimized forward pass for incremental inference

        (i.e., it predicts one time step at a time). If the input order changes

        between time steps, call reorder_incremental_state to update the

        relevant buffers. To generate a fresh sequence, first call

        clear_incremental_state.

        Usage:

        ```

        with model.decoder.incremental_inference():

            for step in range(maxlen):

                out, _ = model.decoder(tokens[:, :step], encoder_out)

                probs = model.get_normalized_probs(out[:, -1, :], log_probs=False)

        ```

        """

        class IncrementalInference(object):

            def __init__(self, decoder):

                self.decoder = decoder

            def __enter__(self):

                self.decoder.incremental_eval(True)

            def __exit__(self, *args):

                self.decoder.incremental_eval(False)

        return IncrementalInference(self)

    def incremental_eval(self, mode=True):

        """Sets the decoder and all children in incremental evaluation mode."""

        assert self._is_incremental_eval != mode, \

            'incremental_eval already set to mode {}'.format(mode)

        self._is_incremental_eval = mode

        if mode:

            self.clear_incremental_state()

        def apply_incremental_eval(module):

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

                module.incremental_eval(mode)

        self.apply(apply_incremental_eval)

    def get_incremental_state(self, key):

        """Return cached state or None if not in incremental inference mode."""

        if self._is_incremental_eval and key in self._incremental_state:

            return self._incremental_state[key]

        return None

    def set_incremental_state(self, key, value):

        """Cache state needed for incremental inference mode."""

        if self._is_incremental_eval:

            self._incremental_state[key] = value

        return value

    def clear_incremental_state(self):

        """Clear all state used for incremental generation.

        **For incremental inference only**

        This should be called before generating a fresh sequence.

        beam_size is required if using BeamableMM.

        """

        if self._is_incremental_eval:

            del self._incremental_state

            self._incremental_state = {}

            def apply_clear_incremental_state(module):

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

                    module.clear_incremental_state()

            self.apply(apply_clear_incremental_state)

    def reorder_incremental_state(self, new_order):

        """Reorder buffered internal state (for incremental generation).

        **For incremental inference only**

    def reorder_incremental_state(self, incremental_state, new_order):

        """Reorder incremental state.

        This should be called when the order of the input has changed from the

        previous time step. A typical use case is beam search, where the input

        order changes between time steps based on the choice of beams.

        order changes between time steps based on the selection of beams.

        """

        if self._is_incremental_eval:

        def apply_reorder_incremental_state(module):

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

                    module.reorder_incremental_state(new_order)

                module.reorder_incremental_state(incremental_state, new_order)

        self.apply(apply_reorder_incremental_state)

    def set_beam_size(self, beam_size):

import torch.nn as nn

import torch.nn.functional as F

from fairseq import utils

from fairseq.data import LanguagePairDataset

from fairseq.modules import BeamableMM, GradMultiply, LearnedPositionalEmbedding, LinearizedConvolution

        else:

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

    def forward(self, prev_output_tokens, encoder_out):

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

        # split and transpose encoder outputs

        encoder_a, encoder_b = self._split_encoder_out(encoder_out)

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

        # embed positions

        positions = self.embed_positions(prev_output_tokens)

        if self._is_incremental_eval:

            # keep only the last token for incremental forward pass

            prev_output_tokens = prev_output_tokens[:, -1:]

        # embed tokens and positions

        x = self.embed_tokens(prev_output_tokens) + positions

        # embed tokens and combine with positional embeddings

        x = self._embed_tokens(prev_output_tokens, incremental_state)

        x += self.embed_positions(prev_output_tokens, incremental_state)

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

        target_embedding = x

        x = self.fc1(x)

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

        x = self._transpose_unless_incremental_eval(x)

        x = self._transpose_if_training(x, incremental_state)

        # temporal convolutions

        avg_attn_scores = None

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

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

            x = conv(x)

            x = conv(x, incremental_state)

            if incremental_state is None:

                x = conv.remove_future_timesteps(x)

            x = F.glu(x, dim=2)

            # attention

            if attention is not None:

                x = self._transpose_unless_incremental_eval(x)

                x = self._transpose_if_training(x, incremental_state)

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

                attn_scores = attn_scores / num_attn_layers

                else:

                    avg_attn_scores.add_(attn_scores)

                x = self._transpose_unless_incremental_eval(x)

                x = self._transpose_if_training(x, incremental_state)

            # residual

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

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

        x = self._transpose_unless_incremental_eval(x)

        x = self._transpose_if_training(x, incremental_state)

        # project back to size of vocabulary

        x = self.fc2(x)

        return x, avg_attn_scores

    def reorder_incremental_state(self, new_order):

        """Reorder buffered internal state (for incremental generation)."""

        super().reorder_incremental_state(new_order)

    def max_positions(self):

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

        return self.embed_positions.max_positions()

            state_dict['decoder.version'] = torch.Tensor([1])

        return state_dict

    def _split_encoder_out(self, encoder_out):

    def _embed_tokens(self, tokens, incremental_state):

        if incremental_state is not None:

            # keep only the last token for incremental forward pass

            tokens = tokens[:, -1:]

        return self.embed_tokens(tokens)

    def _split_encoder_out(self, encoder_out, incremental_state):

        """Split and transpose encoder outputs.

        This is cached when doing incremental inference.

        """

        cached_result = self.get_incremental_state('encoder_out')

        if cached_result:

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

        if cached_result is not None:

            return cached_result

        # transpose only once to speed up attention layers

        encoder_a = encoder_a.transpose(1, 2).contiguous()

        result = (encoder_a, encoder_b)

        return self.set_incremental_state('encoder_out', result)

        if incremental_state is not None:

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

        return result

    def _transpose_unless_incremental_eval(self, x):

        if self._is_incremental_eval:

    def _transpose_if_training(self, x, incremental_state):

        if incremental_state is None:

            x = x.transpose(0, 1)

        return x

        return x.transpose(0, 1)

def Embedding(num_embeddings, embedding_dim, padding_idx):