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

import math

import torch

from torch.autograd import Variable

import torch.nn.functional as F

from fairseq import utils

from . import FairseqCriterion, register_criterion

class LabelSmoothedNLLLoss(torch.autograd.Function):

    @staticmethod

    def forward(ctx, input, target, eps, padding_idx, weights, reduce=True):

        grad_input = input.new(input.size()).zero_()

        target = target.view(target.size(0), 1)

        grad_input = grad_input.scatter_(grad_input.dim() - 1, target, eps - 1)

        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))

        if padding_idx is not None:

            norm -= 1 if weights is None else weights[padding_idx]

            grad_input.select(grad_input.dim() - 1, padding_idx).fill_(0)

        grad_input = grad_input.add(-eps / norm)

        ctx.grad_input = grad_input

        if reduce:

            return input.new([grad_input.view(-1).dot(input.view(-1))])

        else:

            return grad_input * input

    @staticmethod

    def backward(ctx, grad):

        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')

class LabelSmoothedCrossEntropyCriterion(FairseqCriterion):

        """

        net_output = model(**sample['net_input'])

        lprobs = model.get_normalized_probs(net_output, log_probs=True)

        lprobs = lprobs.view(-1, lprobs.size(-1))

        target = sample['target'].view(-1)

        loss = LabelSmoothedNLLLoss.apply(lprobs, target, self.eps, self.padding_idx, None, reduce)

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

        target = sample['target'].unsqueeze(-1)

        non_pad_mask = target.ne(self.padding_idx)

        nll_loss = -lprobs.gather(dim=-1, index=target)[non_pad_mask]

        smooth_loss = -lprobs.sum(dim=-1, keepdim=True)[non_pad_mask]

        if reduce:

            nll_loss = nll_loss.sum()

            smooth_loss = smooth_loss.sum()

        eps_i = self.eps / lprobs.size(-1)

        loss = (1. - self.eps) * nll_loss + eps_i * smooth_loss

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

        logging_output = {

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

    dataset = LanguageDatasets(src, dst, src_dict, dst_dict)

    # Load dataset from binary files

    def all_splits_exist(src, dst):

    def all_splits_exist(src, dst, lang):

        for split in load_splits:

            filename = '{0}.{1}-{2}.{1}.idx'.format(split, src, dst)

            filename = '{0}.{1}-{2}.{3}.idx'.format(split, src, dst, lang)

            if not os.path.exists(os.path.join(path, filename)):

                return False

        return True

    # infer langcode

    if all_splits_exist(src, dst):

    if all_splits_exist(src, dst, src):

        langcode = '{}-{}'.format(src, dst)

    elif all_splits_exist(dst, src):

    elif all_splits_exist(dst, src, src):

        langcode = '{}-{}'.format(dst, src)

    else:

        raise Exception('Dataset cannot be loaded from path: ' + path)

            if not IndexedInMemoryDataset.exists(src_path):

                break

            target_dataset = None

            if IndexedInMemoryDataset.exists(dst_path):

                target_dataset = IndexedInMemoryDataset(dst_path)

            dataset.splits[prefix] = LanguagePairDataset(

                IndexedInMemoryDataset(src_path),

                IndexedInMemoryDataset(dst_path),

                target_dataset,

                pad_idx=dataset.src_dict.pad(),

                eos_idx=dataset.src_dict.eos(),

            )

    def __getitem__(self, i):

        # subtract 1 for 0-based indexing

        source = self.src[i].long() - 1

        target = self.dst[i].long() - 1

        return {

            'id': i,

            'source': source,

            'target': target,

        }

        res = { 'id': i, 'source': source }

        if self.dst:

            res['target'] = self.dst[i].long() - 1

        return res

    def __len__(self):

        return len(self.src)

    def collater(self, samples):

        return LanguagePairDataset.collate(samples, self.pad_idx, self.eos_idx)

        return LanguagePairDataset.collate(samples, self.pad_idx, self.eos_idx, self.dst is not None)

    @staticmethod

    def collate(samples, pad_idx, eos_idx):

    def collate(samples, pad_idx, eos_idx, has_target=True):

        if len(samples) == 0:

            return {}

        id = torch.LongTensor([s['id'] for s in samples])

        src_tokens = merge('source', left_pad=LanguagePairDataset.LEFT_PAD_SOURCE)

        # sort by descending source length

        src_lengths = torch.LongTensor([s['source'].numel() for s in samples])

        src_lengths, sort_order = src_lengths.sort(descending=True)

        id = id.index_select(0, sort_order)

        src_tokens = src_tokens.index_select(0, sort_order)

        prev_output_tokens = None

        target = None

        ntokens = None

        if has_target:

            target = merge('target', left_pad=LanguagePairDataset.LEFT_PAD_TARGET)

            # we create a shifted version of targets for feeding the

            # previous output token(s) into the next decoder step

                left_pad=LanguagePairDataset.LEFT_PAD_TARGET,

                move_eos_to_beginning=True,

            )

        # sort by descending source length

        src_lengths = torch.LongTensor([s['source'].numel() for s in samples])

        src_lengths, sort_order = src_lengths.sort(descending=True)

        id = id.index_select(0, sort_order)

        src_tokens = src_tokens.index_select(0, sort_order)

            prev_output_tokens = prev_output_tokens.index_select(0, sort_order)

            target = target.index_select(0, sort_order)

            ntokens = sum(len(s['target']) for s in samples)

        return {

            'id': id,

            'ntokens': sum(len(s['target']) for s in samples),

            'ntokens': ntokens,

            'net_input': {

                'src_tokens': src_tokens,

                'src_lengths': src_lengths,

    sample_len = 0

    ignored = []

    for idx in map(int, indices):

        if not _valid_size(src.sizes[idx], dst.sizes[idx], max_positions):

        src_size = src.sizes[idx]

        dst_size = dst.sizes[idx] if dst else src_size

        if not _valid_size(src_size, dst_size, max_positions):

            if ignore_invalid_inputs:

                ignored.append(idx)

                continue

            raise Exception((

                "Sample #{} has size (src={}, dst={}) but max size is {}."

                " Skip this example with --skip-invalid-size-inputs-valid-test"

            ).format(idx, src.sizes[idx], dst.sizes[idx], max_positions))

            ).format(idx, src_size, dst_size, max_positions))

        sample_len = max(sample_len, src.sizes[idx], dst.sizes[idx])

        sample_len = max(sample_len, src_size, dst_size)

        num_tokens = (len(batch) + 1) * sample_len

        if yield_batch(idx, num_tokens):

            yield batch

            batch = []

            sample_len = max(src.sizes[idx], dst.sizes[idx])

            sample_len = max(src_size, dst_size)

        batch.append(idx)

                    descending=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 isinstance(dst, IndexedDataset)

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

    if max_tokens is None:

        max_tokens = float('Inf')

    if max_sentences is None:

                    prefix_tokens=s['target'][:, :prefix_size] if prefix_size > 0 else None,

                )

            if timer is not None:

                timer.stop(s['ntokens'])

                timer.stop(sum([len(h[0]['tokens']) for h in hypos]))

            for i, id in enumerate(s['id'].data):

                src = input['src_tokens'].data[i, :]

                # remove padding from ref

                ref = utils.strip_pad(s['target'].data[i, :], self.pad)

                ref = utils.strip_pad(s['target'].data[i, :], self.pad) if s['target'] is not None else None

                yield id, src, ref, hypos[i]

    def generate(self, src_tokens, src_lengths, beam_size=None, maxlen=None, prefix_tokens=None):

                    'alignment': alignment,

                    'positional_scores': pos_scores_i,

                }]

                # return results in the same format as SequenceGenenerator

                # return results in the same format as SequenceGenerator

                yield id, src, ref, hypos

    def score(self, sample):

    # Generate and compute BLEU score

    scorer = bleu.Scorer(dataset.dst_dict.pad(), dataset.dst_dict.eos(), dataset.dst_dict.unk())

    num_sentences = 0

    has_target = True

    with progress_bar.build_progress_bar(args, itr) as t:

        if args.score_reference:

            translations = translator.score_batched_itr(t, cuda=use_cuda, timer=gen_timer)

        wps_meter = TimeMeter()

        for sample_id, src_tokens, target_tokens, hypos in translations:

            # Process input and ground truth

            target_tokens = target_tokens.int().cpu()

            has_target = target_tokens is not None

            target_tokens = target_tokens.int().cpu() if has_target else None

            # Either retrieve the original sentences or regenerate them from tokens.

            if align_dict is not None:

                src_str = dataset.splits[args.gen_subset].src.get_original_text(sample_id)

                target_str = dataset.splits[args.gen_subset].dst.get_original_text(sample_id)

            else:

                src_str = dataset.src_dict.string(src_tokens, args.remove_bpe)

                target_str = dataset.dst_dict.string(target_tokens, args.remove_bpe, escape_unk=True)

                target_str = dataset.dst_dict.string(target_tokens,

                                                     args.remove_bpe,

                                                     escape_unk=True) if has_target else ''

            if not args.quiet:

                print('S-{}\t{}'.format(sample_id, src_str))

                if has_target:

                    print('T-{}\t{}'.format(sample_id, target_str))

            # Process top predictions

                    ))

                # Score only the top hypothesis

                if i == 0:

                if has_target and i == 0:

                    if align_dict is not None or args.remove_bpe is not None:

                        # Convert back to tokens for evaluation with unk replacement and/or without BPE

                        target_tokens = tokenizer.Tokenizer.tokenize(

    print('| Translated {} sentences ({} tokens) in {:.1f}s ({:.2f} tokens/s)'.format(

        num_sentences, gen_timer.n, gen_timer.sum, 1. / gen_timer.avg))

    if has_target:

        print('| Generate {} with beam={}: {}'.format(args.gen_subset, args.beam, scorer.result_string()))