
"""simple docstring"""
import unittest

from transformers import BertGenerationConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import BertGenerationDecoder, BertGenerationEncoder






class      _UpperCAmelCase				:
    def __init__(    self						,    snake_case_						,    snake_case_=13						,    snake_case_=7						,    snake_case_=True						,    snake_case_=True						,    snake_case_=99						,    snake_case_=32						,    snake_case_=5						,    snake_case_=4						,    snake_case_=37						,    snake_case_="gelu"						,    snake_case_=0.1						,    snake_case_=0.1						,    snake_case_=50						,    snake_case_=0.02						,    snake_case_=True						,    snake_case_=None						,    ):
      _snake_case     :						Optional[int]						=   parent
      _snake_case     :						Any						=   batch_size
      _snake_case     :						Union[str, Any]						=   seq_length
      _snake_case     :						List[str]						=   is_training
      _snake_case     :						Any						=   use_input_mask
      _snake_case     :						Optional[Any]						=   vocab_size
      _snake_case     :						str						=   hidden_size
      _snake_case     :						Any						=   num_hidden_layers
      _snake_case     :						Optional[int]						=   num_attention_heads
      _snake_case     :						Tuple						=   intermediate_size
      _snake_case     :						int						=   hidden_act
      _snake_case     :						int						=   hidden_dropout_prob
      _snake_case     :						Any						=   attention_probs_dropout_prob
      _snake_case     :						Tuple						=   max_position_embeddings
      _snake_case     :						List[str]						=   initializer_range
      _snake_case     :						Tuple						=   use_labels
      _snake_case     :						int						=   scope
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Tuple						=   ids_tensor([self.batch_size, self.seq_length]						,    self.vocab_size     )

      _snake_case     :						Any						=   None
      if self.use_input_mask:
        _snake_case     :						Tuple						=   random_attention_mask([self.batch_size, self.seq_length]     )

      if self.use_labels:
        _snake_case     :						str						=   ids_tensor([self.batch_size, self.seq_length]						,    self.vocab_size     )

      _snake_case     :						Optional[Any]						=   self.get_config()

      return config, input_ids, input_mask, token_labels
    def 					lowerCamelCase__	(    self     ):
      return BertGenerationConfig(
          vocab_size=self.vocab_size						,    hidden_size=self.hidden_size						,    num_hidden_layers=self.num_hidden_layers						,    num_attention_heads=self.num_attention_heads						,    intermediate_size=self.intermediate_size						,    hidden_act=self.hidden_act						,    hidden_dropout_prob=self.hidden_dropout_prob						,    attention_probs_dropout_prob=self.attention_probs_dropout_prob						,    max_position_embeddings=self.max_position_embeddings						,    is_decoder=snake_case_						,    initializer_range=self.initializer_range						,    )
    def 					lowerCamelCase__	(    self     ):
      (
          _snake_case					
      )     :						int						=   self.prepare_config_and_inputs()

      _snake_case     :						Union[str, Any]						=   True
      _snake_case     :						Dict						=   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
      _snake_case     :						List[str]						=   ids_tensor([self.batch_size, self.seq_length]						,    vocab_size=2     )

      return (
          config,
          input_ids,
          input_mask,
          token_labels,
          encoder_hidden_states,
          encoder_attention_mask,
      )
    def 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    **snake_case_						,    ):
      _snake_case     :						Tuple						=   BertGenerationEncoder(config=snake_case_     )
      model.to(snake_case_     )
      model.eval()
      _snake_case     :						int						=   model(snake_case_						,    attention_mask=snake_case_     )
      _snake_case     :						List[Any]						=   model(snake_case_     )
      self.parent.assertEqual(result.last_hidden_state.shape						,    (self.batch_size, self.seq_length, self.hidden_size)     )
    def 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    **snake_case_						,    ):
      _snake_case     :						Any						=   True
      _snake_case     :						int						=   BertGenerationEncoder(config=snake_case_     )
      model.to(snake_case_     )
      model.eval()
      _snake_case     :						int						=   model(
          snake_case_						,    attention_mask=snake_case_						,    encoder_hidden_states=snake_case_						,    encoder_attention_mask=snake_case_						,    )
      _snake_case     :						int						=   model(
          snake_case_						,    attention_mask=snake_case_						,    encoder_hidden_states=snake_case_						,    )
      self.parent.assertEqual(result.last_hidden_state.shape						,    (self.batch_size, self.seq_length, self.hidden_size)     )
    def 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    **snake_case_						,    ):
      _snake_case     :						Tuple						=   True
      _snake_case     :						Optional[Any]						=   True
      _snake_case     :						Optional[int]						=   BertGenerationDecoder(config=snake_case_     ).to(snake_case_     ).eval()

      # first forward pass
      _snake_case     :						Dict						=   model(
          snake_case_						,    attention_mask=snake_case_						,    encoder_hidden_states=snake_case_						,    encoder_attention_mask=snake_case_						,    use_cache=snake_case_						,    )
      _snake_case     :						Dict						=   outputs.past_key_values

      # create hypothetical multiple next token and extent to next_input_ids
      _snake_case     :						Optional[int]						=   ids_tensor((self.batch_size, 3)						,    config.vocab_size     )
      _snake_case     :						List[str]						=   ids_tensor((self.batch_size, 3)						,    vocab_size=2     )

      # append to next input_ids and
      _snake_case     :						Optional[int]						=   torch.cat([input_ids, next_tokens]						,    dim=-1     )
      _snake_case     :						List[Any]						=   torch.cat([input_mask, next_mask]						,    dim=-1     )

      _snake_case     :						Optional[int]						=   model(
          snake_case_						,    attention_mask=snake_case_						,    encoder_hidden_states=snake_case_						,    encoder_attention_mask=snake_case_						,    output_hidden_states=snake_case_						,    )["hidden_states"][0]
      _snake_case     :						Tuple						=   model(
          snake_case_						,    attention_mask=snake_case_						,    encoder_hidden_states=snake_case_						,    encoder_attention_mask=snake_case_						,    past_key_values=snake_case_						,    output_hidden_states=snake_case_						,    )["hidden_states"][0]

      # select random slice
      _snake_case     :						int						=   ids_tensor((1,)						,    output_from_past.shape[-1]     ).item()
      _snake_case     :						List[str]						=   output_from_no_past[:, -3:, random_slice_idx].detach()
      _snake_case     :						int						=   output_from_past[:, :, random_slice_idx].detach()

      self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]     )

      # test that outputs are equal for slice
      self.parent.assertTrue(torch.allclose(snake_case_						,    snake_case_						,    atol=1E-3     )     )
    def 					lowerCamelCase__	(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    *snake_case_						,    ):
      _snake_case     :						List[str]						=   BertGenerationDecoder(snake_case_     )
      model.to(snake_case_     )
      model.eval()
      _snake_case     :						Dict						=   model(snake_case_						,    attention_mask=snake_case_						,    labels=snake_case_     )
      self.parent.assertEqual(result.logits.shape						,    (self.batch_size, self.seq_length, self.vocab_size)     )


    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Optional[Any]						=   self.prepare_config_and_inputs()
      _snake_case     :						Tuple						=   {"input_ids": input_ids, "attention_mask": input_mask}
      return config, inputs_dict






@require_torch
class      _UpperCAmelCase				(					_snake_case		,					_snake_case		,					_snake_case		,					unittest.TestCase):
    __lowercase							: Any  =					(BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
    __lowercase							: List[Any]  =					(BertGenerationDecoder,) if is_torch_available() else ()
    __lowercase							: List[str]  =					(
        {"""feature-extraction""": BertGenerationEncoder, """text-generation""": BertGenerationDecoder}
        if is_torch_available()
        else {}
    )
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Optional[Any]						=   BertGenerationEncoderTester(self     )
      _snake_case     :						Optional[int]						=   ConfigTester(self						,    config_class=snake_case_						,    hidden_size=37     )
    def 					lowerCamelCase__	(    self     ):
      self.config_tester.run_common_tests()
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						List[Any]						=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*snake_case_     )
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Any						=   self.model_tester.prepare_config_and_inputs()
      _snake_case     :						Optional[Any]						=   "bert"
      self.model_tester.create_and_check_model(snake_case_						,    snake_case_						,    snake_case_						,    snake_case_     )
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Optional[int]						=   self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_model_as_decoder(*snake_case_     )
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						List[Any]						=   self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_decoder_model_past_large_inputs(*snake_case_     )
    def 					lowerCamelCase__	(    self     ):
      # This regression test was failing with PyTorch < 1.3
      (
          _snake_case					
      )     :						Tuple						=   self.model_tester.prepare_config_and_inputs_for_decoder()

      _snake_case     :						List[str]						=   None

      self.model_tester.create_and_check_model_as_decoder(
          snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_						,    )
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						Optional[Any]						=   self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_for_causal_lm(*snake_case_     )


    @slow
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						int						=   BertGenerationEncoder.from_pretrained("google/bert_for_seq_generation_L-24_bbc_encoder"     )
      self.assertIsNotNone(snake_case_     )






@require_torch
class      _UpperCAmelCase				(					unittest.TestCase):
    @slow
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						str						=   BertGenerationEncoder.from_pretrained("google/bert_for_seq_generation_L-24_bbc_encoder"     )
      _snake_case     :						Optional[int]						=   torch.tensor([[1_01, 75_92, 10_10, 20_26, 38_99, 20_03, 1_01_40, 1_02]]     )
      with torch.no_grad():
        _snake_case     :						int						=   model(snake_case_     )[0]
      _snake_case     :						str						=   torch.Size([1, 8, 10_24]     )
      self.assertEqual(output.shape						,    snake_case_     )
      _snake_case     :						Optional[int]						=   torch.tensor(
          [[[0.1775, 0.0083, -0.0321], [1.6002, 0.1287, 0.3912], [2.1473, 0.5791, 0.6066]]]     )
      self.assertTrue(torch.allclose(output[:, :3, :3]						,    snake_case_						,    atol=1E-4     )     )






@require_torch
class      _UpperCAmelCase				(					unittest.TestCase):
    @slow
    def 					lowerCamelCase__	(    self     ):
      _snake_case     :						List[Any]						=   BertGenerationDecoder.from_pretrained("google/bert_for_seq_generation_L-24_bbc_encoder"     )
      _snake_case     :						List[Any]						=   torch.tensor([[1_01, 75_92, 10_10, 20_26, 38_99, 20_03, 1_01_40, 1_02]]     )
      with torch.no_grad():
        _snake_case     :						List[str]						=   model(snake_case_     )[0]
      _snake_case     :						str						=   torch.Size([1, 8, 5_03_58]     )
      self.assertEqual(output.shape						,    snake_case_     )
      _snake_case     :						Optional[int]						=   torch.tensor(
          [[[-0.5788, -2.5994, -3.7054], [0.0438, 4.7997, 1.8795], [1.5862, 6.6409, 4.4638]]]     )
      self.assertTrue(torch.allclose(output[:, :3, :3]						,    snake_case_						,    atol=1E-4     )     )



"""simple docstring"""
import numpy as np
import torch
from torch.utils.data import Dataset, IterableDataset

from ..utils.generic import ModelOutput






class      _UpperCAmelCase				(					_snake_case):
    def __init__(    self						,    snake_case_						,    snake_case_						,    snake_case_     ):
      _snake_case     :						Optional[int]						=   dataset
      _snake_case     :						str						=   process
      _snake_case     :						int						=   params
    def __len__(    self     ):
      return len(self.dataset     )


    def __getitem__(    self						,    snake_case_     ):
      _snake_case     :						Union[str, Any]						=   self.dataset[i]
      _snake_case     :						Optional[Any]						=   self.process(snake_case_						,    **self.params     )
      return processed






class      _UpperCAmelCase				(					_snake_case):
    def __init__(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_=None     ):
      _snake_case     :						Union[str, Any]						=   loader
      _snake_case     :						Tuple						=   infer
      _snake_case     :						List[Any]						=   params
      if loader_batch_size == 1:
        # Let's spare some time by deactivating altogether
        _snake_case     :						int						=   None
      _snake_case     :						int						=   loader_batch_size

      # Internal bookkeeping
      _snake_case     :						Any						=   None
      _snake_case     :						Dict						=   None
    def __len__(    self     ):
      return len(self.loader     )
    def __iter__(    self     ):
      _snake_case     :						int						=   iter(self.loader     )
      return self
    def 					lowerCamelCase__	(    self     ):
      if isinstance(self._loader_batch_data						,    torch.Tensor     ):
        # Batch data is simple tensor, just fetch the slice
        _snake_case     :						List[Any]						=   self._loader_batch_data[self._loader_batch_index]
      else:
        # Batch data is assumed to be BaseModelOutput (or dict)
        _snake_case     :						int						=   {}
        for k, element in self._loader_batch_data.items():
          if isinstance(snake_case_						,    snake_case_     ):
            # Convert ModelOutput to tuple first
            _snake_case     :						Tuple						=   element.to_tuple()
            if isinstance(element[0]						,    torch.Tensor     ):
              _snake_case     :						Union[str, Any]						=   tuple(el[self._loader_batch_index].unsqueeze(0     ) for el in element     )
            elif isinstance(element[0]						,    np.ndarray     ):
              _snake_case     :						int						=   tuple(np.expand_dims(el[self._loader_batch_index]						,    0     ) for el in element     )
            continue
          if k in {"hidden_states", "past_key_values", "attentions"} and isinstance(snake_case_						,    snake_case_     ):
            # Those are stored as lists of tensors so need specific unbatching.
            if isinstance(element[0]						,    torch.Tensor     ):
              _snake_case     :						Tuple						=   tuple(el[self._loader_batch_index].unsqueeze(0     ) for el in element     )
            elif isinstance(element[0]						,    np.ndarray     ):
              _snake_case     :						Tuple						=   tuple(np.expand_dims(el[self._loader_batch_index]						,    0     ) for el in element     )
            continue
          if element is None:
            # This can happen for optional data that get passed around
            _snake_case     :						Tuple						=   None
          elif isinstance(element[self._loader_batch_index]						,    torch.Tensor     ):
            # Take correct batch data, but make it looked like batch_size=1
            # For compatibility with other methods within transformers

            _snake_case     :						List[Any]						=   element[self._loader_batch_index].unsqueeze(0     )
          elif isinstance(element[self._loader_batch_index]						,    np.ndarray     ):
            # Take correct batch data, but make it looked like batch_size=1
            # For compatibility with other methods within transformers
            _snake_case     :						Union[str, Any]						=   np.expand_dims(element[self._loader_batch_index]						,    0     )
          else:
            # This is typically a list, so no need to `unsqueeze`.
            _snake_case     :						List[Any]						=   element[self._loader_batch_index]
            # Recreate the element by reusing the original class to make it look
            # batch_size=1
        _snake_case     :						int						=   self._loader_batch_data.__class__(snake_case_     )
      self._loader_batch_index += 1
      return result


    def 					lowerCamelCase__	(    self     ):
      if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
        # We are currently unrolling a batch so we just need to return
        # the current item within a batch
        return self.loader_batch_item()

      # We're out of items within a batch
      _snake_case     :						Tuple						=   next(self.iterator     )
      _snake_case     :						Any						=   self.infer(snake_case_						,    **self.params     )
      # We now have a batch of "inferred things".
      if self.loader_batch_size is not None:
        # Try to infer the size of the batch
        if isinstance(snake_case_						,    torch.Tensor     ):
          _snake_case     :						Union[str, Any]						=   processed
        else:
          _snake_case     :						Optional[int]						=   list(processed.keys()     )[0]
          _snake_case     :						List[str]						=   processed[key]
        if isinstance(snake_case_						,    snake_case_     ):
          _snake_case     :						Dict						=   len(snake_case_     )
        else:
          _snake_case     :						Optional[int]						=   first_tensor.shape[0]
        if 0 < observed_batch_size < self.loader_batch_size:
          # could be last batch so we can't unroll as many
          # elements.
          _snake_case     :						Union[str, Any]						=   observed_batch_size
        # Setting internal index to unwrap the batch
        _snake_case     :						str						=   processed
        _snake_case     :						List[Any]						=   0
        return self.loader_batch_item()
      else:
        # We're not unrolling batches
        return processed






class      _UpperCAmelCase				(					_snake_case):
    def __init__(    self						,    snake_case_						,    snake_case_						,    snake_case_						,    snake_case_=None     ):
      super().__init__(snake_case_						,    snake_case_						,    snake_case_     )
    def __iter__(    self     ):
      _snake_case     :						Tuple						=   iter(self.loader     )
      _snake_case     :						List[Any]						=   None
      return self


    def 					lowerCamelCase__	(    self     ):
      if self.subiterator is None:
        _snake_case     :						Optional[Any]						=   self.infer(next(self.iterator     )						,    **self.params     )
      try:
        # Try to return next item
        _snake_case     :						Union[str, Any]						=   next(self.subiterator     )
      except StopIteration:
        # When a preprocess iterator ends, we can start lookig at the next item
        # ChunkIterator will keep feeding until ALL elements of iterator
        # all have created their subiterator and have been iterating against.
        #
        # Another way to look at it, is we're basically flattening lists of lists
        # into a single list, but with generators
        _snake_case     :						str						=   self.infer(next(self.iterator     )						,    **self.params     )
        _snake_case     :						Tuple						=   next(self.subiterator     )
      return processed






class      _UpperCAmelCase				(					_snake_case):
    def __iter__(    self     ):
      _snake_case     :						Optional[Any]						=   iter(self.loader     )
      return self


    def 					lowerCamelCase__	(    self     ):
      # Extremely similar to PipelineIterator in its unpacking mechanism
      # BUT, we have an extra required item which is the presence of `is_last`
      # That is because everything is flattened by `PipelineChunkIterator` we
      # need to keep track of how to regroup here in the original `process`
      # boundaries so that `process` and `postprocess` see the same data.

      # This iterator accumulates items (possibly while unbatching) until it
      # its a `is_last` and then just passes it on to the caller.
      _snake_case     :						Optional[Any]						=   False
      _snake_case     :						Tuple						=   []
      if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
        while self._loader_batch_index < self.loader_batch_size:
          _snake_case     :						Union[str, Any]						=   self.loader_batch_item()
          _snake_case     :						str						=   item.pop("is_last"     )
          accumulator.append(snake_case_     )
          if is_last:
            return accumulator

      while not is_last:
        _snake_case     :						List[str]						=   self.infer(next(self.iterator     )						,    **self.params     )
        if self.loader_batch_size is not None:
          if isinstance(snake_case_						,    torch.Tensor     ):
            _snake_case     :						Union[str, Any]						=   processed
          else:
            _snake_case     :						Tuple						=   list(processed.keys()     )[0]
            _snake_case     :						Tuple						=   processed[key]
          if isinstance(snake_case_						,    snake_case_     ):
            _snake_case     :						Any						=   len(snake_case_     )
          else:
            _snake_case     :						List[Any]						=   first_tensor.shape[0]
          if 0 < observed_batch_size < self.loader_batch_size:
            # could be last batch so we can't unroll as many
            # elements.
            _snake_case     :						Dict						=   observed_batch_size
          _snake_case     :						List[Any]						=   processed
          _snake_case     :						List[str]						=   0
          while self._loader_batch_index < self.loader_batch_size:
            _snake_case     :						Union[str, Any]						=   self.loader_batch_item()
            _snake_case     :						int						=   item.pop("is_last"     )
            accumulator.append(snake_case_     )
            if is_last:
              return accumulator
        else:
          _snake_case     :						Dict						=   processed
          _snake_case     :						Dict						=   item.pop("is_last"     )
          accumulator.append(snake_case_     )
      return accumulator






class      _UpperCAmelCase				(					_snake_case):
    def __init__(    self						,    snake_case_						,    snake_case_     ):
      _snake_case     :						str						=   dataset
      _snake_case     :						Any						=   key
    def __len__(    self     ):
      return len(self.dataset     )


    def __getitem__(    self						,    snake_case_     ):
      return self.dataset[i][self.key]






class      _UpperCAmelCase				(					_snake_case):
    def __init__(    self						,    snake_case_						,    snake_case_						,    snake_case_     ):
      _snake_case     :						int						=   dataset
      _snake_case     :						Any						=   keya
      _snake_case     :						int						=   keya
    def __len__(    self     ):
      return len(self.dataset     )


    def __getitem__(    self						,    snake_case_     ):
      return {"text": self.dataset[i][self.keya], "text_pair": self.dataset[i][self.keya]}

