code: # 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 copy
  import logging
  
  import numpy as np
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  from typing import List
  
  from fairseq import utils
  from fairseq.data import encoders
  
  
  logger = logging.getLogger(__name__)
  
  
  class BARTHubInterface(nn.Module):
      """A simple PyTorch Hub interface to BART.
  
      Usage: https://github.com/pytorch/fairseq/tree/master/examples/BART
      """
  
      def __init__(self, args, task, model):
          super().__init__()
          self.args = args
          self.task = task
          self.model = model
  
          self.bpe = encoders.build_bpe(args)
  
          self.max_positions = min(utils.resolve_max_positions(
              self.task.max_positions(),
              self.model.max_positions(),
          ))
  
          # this is useful for determining the device
          self.register_buffer('_float_tensor', torch.tensor([0], dtype=torch.float))
  
      @property
      def device(self):
          return self._float_tensor.device
  
      def encode(self, sentence: str, *addl_sentences, no_separator=True) -> torch.LongTensor:
          """
          BPE-encode a sentence (or multiple sentences).
  
          Every sequence begins with a beginning-of-sentence (`<s>`) symbol.
          Every sentence ends with an end-of-sentence (`</s>`).
  
          Example (single sentence): `<s> a b c </s>`
          Example (sentence pair): `<s> d e f </s> 1 2 3 </s>`
  
          The BPE encoding follows GPT-2. One subtle detail is that the GPT-2 BPE
          requires leading spaces. For example::
  
              >>> bart.encode('Hello world').tolist()
              [0, 31414, 232, 2]
              >>> bart.encode(' world').tolist()
              [0, 232, 2]
              >>> bart.encode('world').tolist()
              [0, 8331, 2]
          """
          tokens = self.bpe.encode(sentence)
          if len(tokens.split(' ')) > self.max_positions - 2:
              tokens = ' '.join(tokens.split(' ')[:self.max_positions - 2])
          bpe_sentence = '<s> ' + tokens + ' </s>'
          for s in addl_sentences:
              bpe_sentence += (' </s>' if not no_separator else '')
              bpe_sentence += ' ' + self.bpe.encode(s) + ' </s>'
          tokens = self.task.source_dictionary.encode_line(bpe_sentence, append_eos=False)
          return tokens.long()
  
      def decode(self, tokens: torch.LongTensor):
          assert tokens.dim() == 1
          tokens = tokens.cpu().numpy()
          if tokens[0] == self.task.source_dictionary.bos():
              tokens = tokens[1:]  # remove <s>
          eos_mask = (tokens == self.task.source_dictionary.eos())
          doc_mask = eos_mask[1:] & eos_mask[:-1]
          sentences = np.split(tokens, doc_mask.nonzero()[0] + 1)
          sentences = [self.bpe.decode(self.task.source_dictionary.string(s)) for s in sentences]
          if len(sentences) == 1:
              return sentences[0]
          return sentences
  
      def _build_sample(self, src_tokens: List[torch.LongTensor]):
          # assert torch.is_tensor(src_tokens)
          dataset = self.task.build_dataset_for_inference(
              src_tokens,
              [x.numel() for x in src_tokens],
          )
          sample = dataset.collater(dataset)
          sample = utils.apply_to_sample(
              lambda tensor: tensor.to(self.device),
              sample
          )
          return sample
  
      def sample(self, sentences: List[str], beam: int = 1, verbose: bool = False, **kwargs) -> str:
          input = [self.encode(sentence) for sentence in sentences]
          hypos = self.generate(input, beam, verbose, **kwargs)
          return [self.decode(x['tokens']) for x in hypos]
  
      def generate(self, tokens: List[torch.LongTensor], beam: int = 5, verbose: bool = False, **kwargs) -> torch.LongTensor:
          sample = self._build_sample(tokens)
  
          # build generator using current args as well as any kwargs
          gen_args = copy.copy(self.args)
          gen_args.beam = beam
          for k, v in kwargs.items():
              setattr(gen_args, k, v)
          generator = self.task.build_generator([self.model], gen_args)
          translations = self.task.inference_step(
              generator,
              [self.model],
              sample,
              prefix_tokens=sample['net_input']['src_tokens'].new_zeros((len(tokens), 1)).fill_(self.task.source_dictionary.bos()),
          )
  
          if verbose:
              src_str_with_unk = self.string(tokens)
              logger.info('S\t{}'.format(src_str_with_unk))
  
          def getarg(name, default):
              return getattr(gen_args, name, getattr(self.args, name, default))
  
          # Process top predictions
          hypos = [x[0] for x in translations]
          hypos = [v for _, v in sorted(zip(sample['id'].tolist(), hypos))]
          return hypos
  
      def extract_features(self, tokens: torch.LongTensor, return_all_hiddens: bool = False) -> torch.Tensor:
          if tokens.dim() == 1:
              tokens = tokens.unsqueeze(0)
          if tokens.size(-1) > min(self.model.max_positions()):
              raise ValueError('tokens exceeds maximum length: {} > {}'.format(
                  tokens.size(-1), self.model.max_positions()
              ))
          tokens.to(device=self.device),
          prev_output_tokens = tokens.clone()
  
          prev_output_tokens[:, 0] = tokens.gather(
              1,
              (tokens.ne(self.task.source_dictionary.pad()).sum(dim=1)- 1).unsqueeze(-1),
          ).squeeze()
  
          prev_output_tokens[:, 1:] = tokens[:, :-1]
          features, extra = self.model(
              src_tokens=tokens,
              src_lengths=None,
              prev_output_tokens=prev_output_tokens,
              features_only=True,
              return_all_hiddens=return_all_hiddens,
          )
          if return_all_hiddens:
              # convert from T x B x C -> B x T x C
              inner_states = extra['inner_states']
              return [inner_state.transpose(0, 1) for inner_state in inner_states]
          else:
              return features  # just the last layer's features
  
      def register_classification_head(
          self, name: str, num_classes: int = None, embedding_size: int = None, **kwargs
      ):
          self.model.register_classification_head(
              name, num_classes=num_classes, embedding_size=embedding_size, **kwargs
          )
  
      def predict(self, head: str, tokens: torch.LongTensor, return_logits: bool = False):
          if tokens.dim() == 1:
              tokens = tokens.unsqueeze(0)
          features = self.extract_features(tokens.to(device=self.device))
          sentence_representation = features[
              tokens.eq(self.task.source_dictionary.eos()), :
          ].view(features.size(0), -1, features.size(-1))[:, -1, :]
  
          logits = self.model.classification_heads[head](sentence_representation)
          if return_logits:
              return logits
          return F.log_softmax(logits, dim=-1)
   # 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 numpy as np
  import torch
  import math
  
  from . import data_utils, FairseqDataset
  
  
  def collate(
      samples,
      pad_idx,
      eos_idx,
      vocab,
      left_pad_source=False,
      left_pad_target=False,
      input_feeding=True,
  ):
      assert input_feeding
      if len(samples) == 0:
          return {}
  
      def merge(key, left_pad, move_eos_to_beginning=False):
          return data_utils.collate_tokens(
              [s[key] for s in samples],
              pad_idx,
              eos_idx=None,  # use eos_idx of each sample instead of vocab.eos()
              left_pad=left_pad,
              move_eos_to_beginning=move_eos_to_beginning,
          )
  
      id = torch.LongTensor([s['id'] for s in samples])
      src_tokens = merge('source', left_pad=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
      if samples[0].get('target', None) is not None:
          target = merge('target', left_pad=left_pad_target)
          target = target.index_select(0, sort_order)
          ntokens = sum(len(s['target']) for s in samples)
  
          if input_feeding:
              # we create a shifted version of targets for feeding the
              # previous output token(s) into the next decoder step
              prev_output_tokens = merge(
                  'target',
                  left_pad=left_pad_target,
                  move_eos_to_beginning=True,
              )
              prev_output_tokens = prev_output_tokens.index_select(0, sort_order)
      else:
          ntokens = sum(len(s['source']) for s in samples)
  
      batch = {
          'id': id,
          'ntokens': ntokens,
          'net_input': {
              'src_tokens': src_tokens,
              'src_lengths': src_lengths,
          },
          'target': target,
          'nsentences': samples[0]['source'].size(0),
      }
      if prev_output_tokens is not None:
          batch['net_input']['prev_output_tokens'] = prev_output_tokens
  
      return batch
  
  
  class DenoisingDataset(FairseqDataset):
      """
      A wrapper around TokenBlockDataset for BART dataset.
  
      Args:
          dataset (TokenBlockDataset): dataset to wrap
          sizes (List[int]): sentence lengths
          vocab (~fairseq.data.Dictionary): vocabulary
          mask_idx (int): dictionary index used for masked token
          mask_whole_words: only mask whole words. This should be a byte mask
              over vocab indices, indicating whether it is the beginning of a
              word. We will extend any mask to encompass the whole word.
          shuffle (bool, optional): shuffle the elements before batching.
            Default: ``True``
          seed: Seed for random number generator for reproducibility.
          args: argparse arguments.
      """
  
      def __init__(
          self,
          dataset,
          sizes,
          vocab,
          mask_idx,
          mask_whole_words,
          shuffle,
          seed,
          args,
          eos=None
      ):
          self.dataset = dataset
  
          self.sizes = sizes
  
          self.vocab = vocab
          self.shuffle = shuffle
          self.seed = seed
          self.mask_idx = mask_idx
          self.mask_whole_word = mask_whole_words
          self.mask_ratio = args.mask
          self.random_ratio = args.mask_random
          self.insert_ratio = args.insert
          self.rotate_ratio = args.rotate
          self.permute_sentence_ratio = args.permute_sentences
          self.eos = (eos if eos is not None else vocab.eos())
  
          if args.bpe != 'gpt2':
              self.full_stop_index = self.vocab.eos()
          else:
              assert args.bpe == 'gpt2'
              self.full_stop_index = self.vocab.index('13')
  
          self.replace_length = args.replace_length
          if not self.replace_length in [-1, 0, 1]:
              raise ValueError(f'invalid arg: replace_length={self.replace_length}')
          if not args.mask_length in ['subword', 'word', 'span-poisson']:
              raise ValueError(f'invalid arg: mask-length={args.mask_length}')
          if args.mask_length == 'subword' and not args.replace_length in [0, 1]:
              raise ValueError(f'if using subwords, use replace-length=1 or 0')
  
          self.mask_span_distribution = None
          if args.mask_length == 'span-poisson':
              _lambda = args.poisson_lambda
  
              lambda_to_the_k = 1
              e_to_the_minus_lambda = math.exp(-_lambda)
              k_factorial = 1
              ps = []
              for k in range(0, 128):
                  ps.append(e_to_the_minus_lambda * lambda_to_the_k / k_factorial)
                  lambda_to_the_k *= _lambda
                  k_factorial *= (k + 1)
                  if ps[-1] < 0.0000001:
                      break
              ps = torch.FloatTensor(ps)
              self.mask_span_distribution = torch.distributions.Categorical(ps)
  
          self.epoch = 0
  
      def set_epoch(self, epoch, **unused):
          self.epoch = epoch
  
      def __getitem__(self, index):
          with data_utils.numpy_seed(self.seed, self.epoch, index):
              tokens = self.dataset[index]
              assert tokens[-1] == self.eos
              source, target = tokens, tokens.clone()
  
              if self.permute_sentence_ratio > 0.0:
                  source = self.permute_sentences(source, self.permute_sentence_ratio)
  
              if self.mask_ratio > 0:
                  source = self.add_whole_word_mask(source, self.mask_ratio)
  
              if self.insert_ratio > 0:
                  source = self.add_insertion_noise(source, self.insert_ratio)
  
              if self.rotate_ratio > 0.0 and np.random.random() < self.rotate_ratio:
                  source = self.add_rolling_noise(source)
  
          assert (source >= 0).all()
          assert (source[1:-1] >= 1).all()
          assert (source <= len(self.vocab)).all()
          assert source[0] == self.vocab.bos()
          assert source[-1] == self.eos
          return {
              'id': index,
              'source': source,
              'target': target,
          }
  
      def __len__(self):
          return len(self.dataset)
  
      def permute_sentences(self, source, p=1.0):
          full_stops = (source == self.full_stop_index)
          # Pretend it ends with a full stop so last span is a sentence
          full_stops[-2] = 1
  
          # Tokens that are full stops, where the previous token is not
          sentence_ends = (full_stops[1:] * ~full_stops[:-1]).nonzero() + 2
          result = source.clone()
  
          num_sentences = sentence_ends.size(0)
          num_to_permute = math.ceil((num_sentences * 2 * p) / 2.0)
          substitutions = torch.randperm(num_sentences)[:num_to_permute]
          ordering = torch.arange(0, num_sentences)
          ordering[substitutions] = substitutions[torch.randperm(num_to_permute)]
  
          # Ignore <bos> at start
          index = 1
          for i in ordering:
              sentence = source[(sentence_ends[i - 1] if i > 0 else 1):sentence_ends[i]]
              result[index:index + sentence.size(0)] = sentence
              index += sentence.size(0)
          return result
  
      def word_starts(self, source):
          if self.mask_whole_word is not None:
              is_word_start = self.mask_whole_word.gather(0, source)
          else:
              is_word_start = torch.ones(source.size())
          is_word_start[0] = 0
          is_word_start[-1] = 0
          return is_word_start
  
      def add_whole_word_mask(self, source, p):
          is_word_start = self.word_starts(source)
          num_to_mask = int(math.ceil(is_word_start.float().sum() * p))
          num_inserts = 0
          if num_to_mask == 0:
              return source
  
          if self.mask_span_distribution is not None:
              lengths = self.mask_span_distribution.sample(sample_shape=(num_to_mask,))
  
              # Make sure we have enough to mask
              cum_length = torch.cumsum(lengths, 0)
              while cum_length[-1] < num_to_mask:
                  lengths = torch.cat([lengths, self.mask_span_distribution.sample(sample_shape=(num_to_mask,))], dim=0)
                  cum_length = torch.cumsum(lengths, 0)
  
              # Trim to masking budget
              i = 0
              while cum_length[i] < num_to_mask:
                  i += 1
              lengths[i] = num_to_mask - (0 if i == 0 else cum_length[i - 1])
              num_to_mask = i + 1
              lengths = lengths[:num_to_mask]
  
              # Handle 0-length mask (inserts) separately
              lengths = lengths[lengths > 0]
              num_inserts = num_to_mask - lengths.size(0)
              num_to_mask -= num_inserts
              if num_to_mask == 0:
                  return self.add_insertion_noise(source, num_inserts / source.size(0))
  
              assert (lengths > 0).all()
          else:
              lengths = torch.ones((num_to_mask,)).long()
          assert is_word_start[-1] == 0
          word_starts = is_word_start.nonzero()
          indices = word_starts[torch.randperm(word_starts.size(0))[:num_to_mask]].squeeze(1)
          mask_random = torch.FloatTensor(num_to_mask).uniform_() < self.random_ratio
  
          source_length = source.size(0)
          assert source_length - 1 not in indices
          to_keep = torch.ones(source_length, dtype=torch.bool)
          is_word_start[-1] = 255 # acts as a long length, so spans don't go over the end of doc
          if self.replace_length == 0:
              to_keep[indices] = 0
          else:
              # keep index, but replace it with [MASK]
              source[indices] = self.mask_idx
              source[indices[mask_random]] = torch.randint(1, len(self.vocab), size=(mask_random.sum(),))
  
          if self.mask_span_distribution is not None:
              assert len(lengths.size()) == 1
              assert lengths.size() == indices.size()
              lengths -= 1
              while indices.size(0) > 0:
                  assert lengths.size() == indices.size()
                  lengths -= is_word_start[indices + 1].long()
                  uncompleted = lengths >= 0
                  indices = indices[uncompleted] + 1
                  mask_random = mask_random[uncompleted]
                  lengths = lengths[uncompleted]
                  if self.replace_length != -1:
                      # delete token
                      to_keep[indices] = 0
                  else:
                      # keep index, but replace it with [MASK]
                      source[indices] = self.mask_idx
                      source[indices[mask_random]] = torch.randint(1, len(self.vocab), size=(mask_random.sum(),))
          else:
              # A bit faster when all lengths are 1
              while indices.size(0) > 0:
                  uncompleted = is_word_start[indices + 1] == 0
                  indices = indices[uncompleted] + 1
                  mask_random = mask_random[uncompleted]
                  if self.replace_length != -1:
                      # delete token
                      to_keep[indices] = 0
                  else:
                      # keep index, but replace it with [MASK]
                      source[indices] = self.mask_idx
                      source[indices[mask_random]] = torch.randint(1, len(self.vocab), size=(mask_random.sum(),))
  
                  assert source_length - 1 not in indices
  
          source = source[to_keep]
  
          if num_inserts > 0:
              source = self.add_insertion_noise(source, num_inserts / source.size(0))
  
          return source
  
      def add_permuted_noise(self, tokens, p):
          num_words = len(tokens)
          num_to_permute = math.ceil(((num_words * 2) * p) / 2.0)
          substitutions = torch.randperm(num_words - 2)[:num_to_permute] + 1
          tokens[substitutions] = tokens[substitutions[torch.randperm(num_to_permute)]]
          return tokens
  
      def add_rolling_noise(self, tokens):
          offset = np.random.randint(1, max(1, tokens.size(-1) - 1) + 1)
          tokens = torch.cat(
              (tokens[0:1], tokens[offset:-1], tokens[1:offset], tokens[-1:]),
              dim=0,
          )
          return tokens
  
      def add_insertion_noise(self, tokens, p):
          if p == 0.0:
              return tokens
  
          num_tokens = len(tokens)
          n = int(math.ceil(num_tokens * p))
  
          noise_indices = torch.randperm(num_tokens + n - 2)[:n] + 1
          noise_mask = torch.zeros(size=(num_tokens + n,), dtype=torch.bool)
          noise_mask[noise_indices] = 1
          result = torch.LongTensor(n + len(tokens)).fill_(-1)
  
          num_random = int(math.ceil(n * self.random_ratio))
          result[noise_indices[num_random:]] = self.mask_idx
          result[noise_indices[:num_random]] = torch.randint(low=1, high=len(self.vocab), size=(num_random,))
  
          result[~noise_mask] = tokens
  
          assert (result >= 0).all()
          return result
  
      def collater(self, samples):
          """Merge a list of samples to form a mini-batch.
          Args:
              samples (List[dict]): samples to collate
          Returns:
              dict: a mini-batch of data
          """
          return collate(samples, self.vocab.pad(), self.eos, self.vocab)
  
      def num_tokens(self, index):
          """Return the number of tokens in a sample. This value is used to
          enforce ``--max-tokens`` during batching."""
          return self.sizes[index]
  
      def size(self, index):
          """Return an example's size as a float or tuple. This value is used when
          filtering a dataset with ``--max-positions``."""
          return self.sizes[index]
  
      def ordered_indices(self):
          """Return an ordered list of indices. Batches will be constructed based
          on this order."""
          if self.shuffle:
              indices = np.random.permutation(len(self))
          else:
              indices = np.arange(len(self))
          return indices[np.argsort(self.sizes[indices], kind='mergesort')]
  
      def prefetch(self, indices):
          self.src.prefetch(indices)
          self.tgt.prefetch(indices)
  
      @property
      def supports_prefetch(self):
          return (
              hasattr(self.src, 'supports_prefetch')
              and self.src.supports_prefetch
              and hasattr(self.tgt, 'supports_prefetch')
              and self.tgt.supports_prefetch
          )
  
apis: torch.nn.functional.log_softmax torch.tensor
apis: torch.ones torch.FloatTensor torch.distributions.Categorical numpy.argsort numpy.random.random torch.cumsum torch.arange torch.randperm torch.zeros torch.LongTensor torch.cat