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 logging
  import warnings
  import torch
  import numpy as np
  
  from data import data_utils
  from data.ofa_dataset import OFADataset
  
  logger = logging.getLogger(__name__)
  warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
  
  
  def collate(samples, pad_idx, eos_idx):
      if len(samples) == 0:
          return {}
  
      def merge(key):
          return data_utils.collate_tokens(
              [s[key] for s in samples],
              pad_idx,
              eos_idx=eos_idx,
          )
  
      src_tokens = merge("source")
      src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
  
      ref_dict = None
      if samples[0].get("ref_dict", None) is not None:
          ref_dict = np.array([s['ref_dict'] for s in samples])
  
      constraint_masks = None
      if samples[0].get("constraint_mask", None) is not None:
          constraint_masks = merge("constraint_mask")
  
      prev_output_tokens = None
      target = None
      if samples[0].get("target", None) is not None:
          target = merge("target")
          tgt_lengths = torch.LongTensor(
              [s["target"].ne(pad_idx).long().sum() for s in samples]
          )
          ntokens = tgt_lengths.sum().item()
  
          if samples[0].get("prev_output_tokens", None) is not None:
              prev_output_tokens = merge("prev_output_tokens")
      else:
          ntokens = src_lengths.sum().item()
  
      batch = {
          "nsentences": len(samples),
          "ntokens": ntokens,
          "net_input": {
              "src_tokens": src_tokens,
              "src_lengths": src_lengths,
              "prev_output_tokens": prev_output_tokens
          },
          "ref_dict": ref_dict,
          "constraint_masks": constraint_masks,
          "target": target,
      }
  
      return batch
  
  
  class SST2Dataset(OFADataset):
      def __init__(
          self,
          split,
          dataset,
          bpe,
          src_dict,
          tgt_dict=None,
          max_src_length=512,
          max_tgt_length=30,
          constraint_trie=None,
          prompt_type="none"
      ):
          super().__init__(split, dataset, bpe, src_dict, tgt_dict)
          self.max_src_length = max_src_length
          self.max_tgt_length = max_tgt_length
          self.constraint_trie = constraint_trie
          self.prompt_type = prompt_type
  
      def __getitem__(self, index):
          sentence, label = self.dataset[index]
          if label == '0':
              label = 'negative'
          elif label == '1':
              label = 'positive'
          else:
              raise NotImplementedError
  
          sentence = ' '.join(sentence.lower().strip().split()[:self.max_src_length])
          src_item = self.encode_text(' is the sentiment of text " {} " positive or negative?'.format(sentence))
          tgt_item = self.encode_text(" {}".format(label))
          assert tgt_item.size(0) == 1
          ref_dict = {label: 1.0}
  
          src_item = torch.cat([self.bos_item, src_item, self.eos_item])
          if self.prompt_type == 'none':
              prev_output_item = self.bos_item
              target_item = tgt_item
          elif self.prompt_type == 'src':
              prev_output_item = src_item.clone()
              target_item = torch.cat([prev_output_item[1:], tgt_item])
          elif self.prompt_type == 'prev_output':
              prev_output_item = src_item[:-1].clone()
              target_item = torch.cat([prev_output_item[1:], tgt_item])
          else:
              raise NotImplementedError
          target_item[:-1] = self.tgt_dict.pad()
  
          example = {
              "source": src_item,
              "target": target_item,
              "prev_output_tokens": prev_output_item,
              "ref_dict": ref_dict,
          }
          if self.constraint_trie is not None:
              constraint_mask = torch.zeros((len(prev_output_item), len(self.tgt_dict))).bool()
              constraint_nodes = self.constraint_trie.get_next_layer(self.bos_item.tolist())
              constraint_mask[-1][constraint_nodes] = True
              example["constraint_mask"] = constraint_mask
          return example
  
      def collater(self, samples, pad_to_length=None):
          """Merge a list of samples to form a mini-batch.
          Args:
              samples (List[dict]): samples to collate
          Returns:
              dict: a mini-batch with the following keys:
          """
          return collate(samples, pad_idx=self.pad, eos_idx=self.eos)