text: # Copyright (c) Facebook, Inc. and its affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  import os
  import numpy as np
  
  import torch
  
  from collections import OrderedDict
  from fairseq import utils
  from fairseq.data import (
      data_utils,
      Dictionary,
      TokenBlockDataset,
  )
  from fairseq.tasks import FairseqTask, register_task
  from .masked_dataset import MaskedLanguagePairDataset
  from .bert_dictionary import BertDictionary
  
  
  @register_task('masked_s2s')
  class MaskedS2STask(FairseqTask):
      """
      Train a sequence-to-sequence task
  
      Args:
          dictionary (~fairseq.data.Dictionary): the dictionary for the input of
              the language model
      """
  
      @staticmethod
      def add_args(parser):
          """Add task-specific arguments to the parser."""
          # fmt: off
          parser.add_argument('data', help='path to data directory')
          parser.add_argument('--sample-break-mode', default='none',
                              choices=['none', 'complete', 'complete_doc', 'eos'],
                              help='If omitted or "none", fills each sample with tokens-per-sample '
                                   'tokens. If set to "complete", splits samples only at the end '
                                   'of sentence, but may include multiple sentences per sample. '
                                   '"complete_doc" is similar but respects doc boundaries. '
                                   'If set to "eos", includes only one sentence per sample.')
          parser.add_argument('--tokens-per-sample', default=512, type=int,
                              help='max number of tokens per sample for text dataset')
          parser.add_argument('--lazy-load', action='store_true',
                              help='load the dataset lazily')
          parser.add_argument('--raw-text', default=False, action='store_true',
                              help='load raw text dataset')
  
          parser.add_argument('--mask-s2s-prob', default=0.15, type=float,
                              help='probability of replacing a token with mask')
          parser.add_argument('--mask-s2s-mask-keep-rand', default="0.8,0.1,0.1", type=str,
                              help='Word prediction probability for decoder mask')
  
          # fmt: on
  
      def __init__(self, args, dictionary):
          super().__init__(args)
          self.dictionary = dictionary
  
      @classmethod
      def setup_task(cls, args, **kwargs):
          """Setup the task (e.g., load dictionaries).
  
          Args:
              args (argparse.Namespace): parsed command-line arguments
          """
          if getattr(args, 'raw_text', False):
              utils.deprecation_warning('--raw-text is deprecated, please use --dataset-impl=raw')
              args.dataset_impl = 'raw'
          elif getattr(args, 'lazy_load', False):
              utils.deprecation_warning('--lazy-load is deprecated, please use --dataset-impl=lazy')
              args.dataset_impl = 'lazy'
  
          paths = args.data.split(':')
  
          dictionary = cls.load_dictionary(os.path.join(paths[0], 'dict.txt'))
          print('| dictionary: {} types'.format(len(dictionary)))
          return cls(args, dictionary)
  
      @classmethod
      def load_dictionary(cls, filename):
          return BertDictionary.load_from_file(filename)
  
      def train_step(self, sample, model, criterion, optimizer, ignore_grad=False):
          model.train()
          loss, sample_size, logging_output = criterion(model, sample)
          if ignore_grad:
              loss *= 0
          optimizer.backward(loss)
          return loss, sample_size, logging_output
  
      def valid_step(self, sample, model, criterion):
          model.eval()
          with torch.no_grad():
              loss, sample_size, logging_output = criterion(model, sample)
          return loss, sample_size, logging_output
  
      def build_model(self, args):
          from fairseq import models
          model = models.build_model(args, self)
          return model
  
      def load_dataset(self, split, epoch=0, combine=False, **kwargs):
          """Load a given dataset split.
  
          Args:
              split (str): name of the split (e.g., train, valid, test)
          """
          paths = self.args.data.split(':')
          assert len(paths) > 0
          data_path = paths[epoch % len(paths)]
          split_path = os.path.join(data_path, split)
  
          dataset = data_utils.load_indexed_dataset(
              split_path,
              self.dictionary,
              self.args.dataset_impl,
              combine=combine,
          )
          if dataset is None:
              raise FileNotFoundError('Dataset not found: {} ({})'.format(split, split_path))
  
          self.datasets[split] = self.build_s2s_dataset(dataset)
  
      def build_s2s_dataset(self, dataset):
          dataset = TokenBlockDataset(
              dataset,
              dataset.sizes,
              self.args.tokens_per_sample,
              pad=self.source_dictionary.pad(),
              eos=self.source_dictionary.eos(),
              break_mode=self.args.sample_break_mode,
          )
  
          pred_probs = torch.FloatTensor([float(x) for x in self.args.mask_s2s_mask_keep_rand.split(',')])
  
          s2s_dataset = MaskedLanguagePairDataset(
              dataset, dataset.sizes, self.source_dictionary,
              shuffle=True, mask_prob=self.args.mask_s2s_prob,
              pred_probs=pred_probs,
          )
          return s2s_dataset
  
      def build_dataset_for_inference(self, src_tokens, src_lengths):
          raise NotImplementedError
  
      def inference_step(self, generator, models, sample, prefix_tokens=None):
          raise NotImplementedError
  
      @property
      def source_dictionary(self):
          """Return the :class:`~fairseq.data.Dictionary` for the language
          model."""
          return self.dictionary
  
      @property
      def target_dictionary(self):
          """Return the :class:`~fairseq.data.Dictionary` for the language
          model."""
          return self.dictionary
  
      def max_positions(self):
          max_positions = 1024
          if hasattr(self.args, 'max_positions'):
              max_positions = min(max_positions, self.args.max_positions)
          if hasattr(self.args, 'max_source_positions'):
              max_positions = min(max_positions, self.args.max_source_positions)
          if hasattr(self.args, 'max_target_positions'):
              max_positions = min(max_positions, self.args.max_target_positions)
          return (max_positions, max_positions)
  