





'''simple docstring'''

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

from ..utils.generic import ModelOutput






class        UpperCAmelCase     (	UpperCamelCase__    ):
 def __init__(    self				:str						,	lowercase_				:Dict						,	lowercase_				:Any						,	lowercase_				:int							)->							Tuple:
     A__    =      dataset
     A__    =      process
     A__    =      params
 def __len__(    self				:Dict							)->							Any:
     return len(self.dataset							)



 def __getitem__(    self				:Union[str, Any]						,	lowercase_				:int							)->							Optional[int]:
     A__    =      self.dataset[i]
     A__    =      self.process(lowercase_						,	**self.params							)
     return processed






class        UpperCAmelCase     (	UpperCamelCase__    ):
 def __init__(    self				:str						,	lowercase_				:int						,	lowercase_				:List[Any]						,	lowercase_				:Optional[int]						,	lowercase_				:List[Any]=None							)->							int:
     A__    =      loader
     A__    =      infer
     A__    =      params
     if loader_batch_size == 1:
         # Let's spare some time by deactivating altogether
         A__    =      None
     A__    =      loader_batch_size

     # Internal bookkeeping
     A__    =      None
     A__    =      None
 def __len__(    self				:str							)->							Tuple:
     return len(self.loader							)
 def __iter__(    self				:List[str]							)->							str:
     A__    =      iter(self.loader							)
     return self
 def 					UpperCAmelCase_							(    self				:List[Any]							)->							Optional[Any]:
     if isinstance(self._loader_batch_data						,	torch.Tensor							):
         # Batch data is simple tensor, just fetch the slice
         A__    =      self._loader_batch_data[self._loader_batch_index]
     else:
         # Batch data is assumed to be BaseModelOutput (or dict)
         A__    =      {}
         for k, element in self._loader_batch_data.items():
             if isinstance(lowercase_						,	lowercase_							):
                 # Convert ModelOutput to tuple first
                 A__    =      element.to_tuple()
                 if isinstance(element[0]						,	torch.Tensor							):
                     A__    =      tuple(el[self._loader_batch_index].unsqueeze(0							) for el in element							)
                 elif isinstance(element[0]						,	np.ndarray							):
                     A__    =      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(lowercase_						,	lowercase_							):
                 # Those are stored as lists of tensors so need specific unbatching.
                 if isinstance(element[0]						,	torch.Tensor							):
                     A__    =      tuple(el[self._loader_batch_index].unsqueeze(0							) for el in element							)
                 elif isinstance(element[0]						,	np.ndarray							):
                     A__    =      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
                 A__    =      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

                 A__    =      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
                 A__    =      np.expand_dims(element[self._loader_batch_index]						,	0							)
             else:
                 # This is typically a list, so no need to `unsqueeze`.
                 A__    =      element[self._loader_batch_index]
            # Recreate the element by reusing the original class to make it look
            # batch_size=1
         A__    =      self._loader_batch_data.__class__(lowercase_							)
     self._loader_batch_index += 1
     return result



 def 					UpperCAmelCase_							(    self				:Tuple							)->							Dict:
     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
     A__    =      next(self.iterator							)
     A__    =      self.infer(lowercase_						,	**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(lowercase_						,	torch.Tensor							):
             A__    =      processed
         else:
             A__    =      list(processed.keys()							)[0]
             A__    =      processed[key]
         if isinstance(lowercase_						,	lowercase_							):
             A__    =      len(lowercase_							)
         else:
             A__    =      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.
             A__    =      observed_batch_size
         # Setting internal index to unwrap the batch
         A__    =      processed
         A__    =      0
         return self.loader_batch_item()
     else:
         # We're not unrolling batches
         return processed






class        UpperCAmelCase     (	UpperCamelCase__    ):
 def __init__(    self				:List[str]						,	lowercase_				:Any						,	lowercase_				:List[Any]						,	lowercase_				:str						,	lowercase_				:Dict=None							)->							Union[str, Any]:
     super().__init__(lowercase_						,	lowercase_						,	lowercase_							)
 def __iter__(    self				:Optional[int]							)->							Dict:
     A__    =      iter(self.loader							)
     A__    =      None
     return self



 def 					UpperCAmelCase_							(    self				:str							)->							Optional[int]:
     if self.subiterator is None:
         A__    =      self.infer(next(self.iterator							)						,	**self.params							)
     try:
         # Try to return next item
         A__    =      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
         A__    =      self.infer(next(self.iterator							)						,	**self.params							)
         A__    =      next(self.subiterator							)
     return processed






class        UpperCAmelCase     (	UpperCamelCase__    ):
 def __iter__(    self				:Union[str, Any]							)->							List[str]:
     A__    =      iter(self.loader							)
     return self



 def 					UpperCAmelCase_							(    self				:Optional[Any]							)->							Union[str, Any]:
     # 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.
     A__    =      False
     A__    =      []
     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:
             A__    =      self.loader_batch_item()
             A__    =      item.pop("is_last"							)
             accumulator.append(lowercase_							)
             if is_last:
                 return accumulator

     while not is_last:
         A__    =      self.infer(next(self.iterator							)						,	**self.params							)
         if self.loader_batch_size is not None:
             if isinstance(lowercase_						,	torch.Tensor							):
                 A__    =      processed
             else:
                 A__    =      list(processed.keys()							)[0]
                 A__    =      processed[key]
             if isinstance(lowercase_						,	lowercase_							):
                 A__    =      len(lowercase_							)
             else:
                 A__    =      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.
                 A__    =      observed_batch_size
             A__    =      processed
             A__    =      0
             while self._loader_batch_index < self.loader_batch_size:
                 A__    =      self.loader_batch_item()
                 A__    =      item.pop("is_last"							)
                 accumulator.append(lowercase_							)
                 if is_last:
                     return accumulator
         else:
             A__    =      processed
             A__    =      item.pop("is_last"							)
             accumulator.append(lowercase_							)
     return accumulator






class        UpperCAmelCase     (	UpperCamelCase__    ):
 def __init__(    self				:str						,	lowercase_				:Dataset						,	lowercase_				:str							)->							Optional[Any]:
     A__    =      dataset
     A__    =      key
 def __len__(    self				:int							)->							Optional[Any]:
     return len(self.dataset							)



 def __getitem__(    self				:int						,	lowercase_				:List[Any]							)->							Dict:
     return self.dataset[i][self.key]






class        UpperCAmelCase     (	UpperCamelCase__    ):
 def __init__(    self				:str						,	lowercase_				:Dataset						,	lowercase_				:str						,	lowercase_				:str							)->							Union[str, Any]:
     A__    =      dataset
     A__    =      keya
     A__    =      keya
 def __len__(    self				:Union[str, Any]							)->							str:
     return len(self.dataset							)



 def __getitem__(    self				:Optional[Any]						,	lowercase_				:Tuple							)->							List[Any]:
     return {"text": self.dataset[i][self.keya], "text_pair": self.dataset[i][self.keya]}










'''simple docstring'''

from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers.models.distilbert.modeling_tf_distilbert import (
      TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFDistilBertForMaskedLM,
      TFDistilBertForMultipleChoice,
      TFDistilBertForQuestionAnswering,
      TFDistilBertForSequenceClassification,
      TFDistilBertForTokenClassification,
      TFDistilBertModel,
  )






class        UpperCAmelCase     :
 def __init__(    self				:str						,	lowercase_				:str						,	)->							str:
     A__    =      parent
     A__    =      13
     A__    =      7
     A__    =      True
     A__    =      True
     A__    =      False
     A__    =      True
     A__    =      99
     A__    =      32
     A__    =      2
     A__    =      4
     A__    =      37
     A__    =      "gelu"
     A__    =      0.1
     A__    =      0.1
     A__    =      5_12
     A__    =      16
     A__    =      2
     A__    =      0.0_2
     A__    =      3
     A__    =      4
     A__    =      None
 def 					UpperCAmelCase_							(    self				:Union[str, Any]							)->							int:
     A__    =      ids_tensor([self.batch_size, self.seq_length]						,	self.vocab_size							)

     A__    =      None
     if self.use_input_mask:
         A__    =      random_attention_mask([self.batch_size, self.seq_length]							)

     A__    =      None
     A__    =      None
     A__    =      None
     if self.use_labels:
         A__    =      ids_tensor([self.batch_size]						,	self.type_sequence_label_size							)
         A__    =      ids_tensor([self.batch_size, self.seq_length]						,	self.num_labels							)
         A__    =      ids_tensor([self.batch_size]						,	self.num_choices							)

     A__    =      DistilBertConfig(
         vocab_size=self.vocab_size						,	dim=self.hidden_size						,	n_layers=self.num_hidden_layers						,	n_heads=self.num_attention_heads						,	hidden_dim=self.intermediate_size						,	hidden_act=self.hidden_act						,	dropout=self.hidden_dropout_prob						,	attention_dropout=self.attention_probs_dropout_prob						,	max_position_embeddings=self.max_position_embeddings						,	initializer_range=self.initializer_range						,	)

     return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
 def 					UpperCAmelCase_							(    self				:str						,	lowercase_				:Optional[int]						,	lowercase_				:List[str]						,	lowercase_				:Any						,	lowercase_				:Union[str, Any]						,	lowercase_				:Optional[int]						,	lowercase_				:str							)->							List[str]:
     A__    =      TFDistilBertModel(config=lowercase_							)
     A__    =      {"input_ids": input_ids, "attention_mask": input_mask}

     A__    =      model(lowercase_							)

     A__    =      [input_ids, input_mask]

     A__    =      model(lowercase_							)

     self.parent.assertEqual(result.last_hidden_state.shape						,	(self.batch_size, self.seq_length, self.hidden_size)							)
 def 					UpperCAmelCase_							(    self				:List[str]						,	lowercase_				:str						,	lowercase_				:Optional[Any]						,	lowercase_				:Optional[int]						,	lowercase_				:Optional[int]						,	lowercase_				:Optional[int]						,	lowercase_				:Union[str, Any]							)->							Optional[int]:
     A__    =      TFDistilBertForMaskedLM(config=lowercase_							)
     A__    =      {"input_ids": input_ids, "attention_mask": input_mask}
     A__    =      model(lowercase_							)
     self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.seq_length, self.vocab_size)							)
 def 					UpperCAmelCase_							(    self				:Any						,	lowercase_				:str						,	lowercase_				:str						,	lowercase_				:Optional[int]						,	lowercase_				:str						,	lowercase_				:List[Any]						,	lowercase_				:Union[str, Any]							)->							Optional[int]:
     A__    =      TFDistilBertForQuestionAnswering(config=lowercase_							)
     A__    =      {
         "input_ids": input_ids,
         "attention_mask": input_mask,
     }
     A__    =      model(lowercase_							)
     self.parent.assertEqual(result.start_logits.shape						,	(self.batch_size, self.seq_length)							)
     self.parent.assertEqual(result.end_logits.shape						,	(self.batch_size, self.seq_length)							)
 def 					UpperCAmelCase_							(    self				:Union[str, Any]						,	lowercase_				:Optional[int]						,	lowercase_				:Any						,	lowercase_				:Dict						,	lowercase_				:Tuple						,	lowercase_				:Optional[Any]						,	lowercase_				:Optional[int]							)->							Any:
     A__    =      self.num_labels
     A__    =      TFDistilBertForSequenceClassification(lowercase_							)
     A__    =      {"input_ids": input_ids, "attention_mask": input_mask}
     A__    =      model(lowercase_							)
     self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.num_labels)							)
 def 					UpperCAmelCase_							(    self				:str						,	lowercase_				:Optional[Any]						,	lowercase_				:List[Any]						,	lowercase_				:Dict						,	lowercase_				:Tuple						,	lowercase_				:int						,	lowercase_				:Union[str, Any]							)->							str:
     A__    =      self.num_choices
     A__    =      TFDistilBertForMultipleChoice(lowercase_							)
     A__    =      tf.tile(tf.expand_dims(lowercase_						,	1							)						,	(1, self.num_choices, 1)							)
     A__    =      tf.tile(tf.expand_dims(lowercase_						,	1							)						,	(1, self.num_choices, 1)							)
     A__    =      {
         "input_ids": multiple_choice_inputs_ids,
         "attention_mask": multiple_choice_input_mask,
     }
     A__    =      model(lowercase_							)
     self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.num_choices)							)
 def 					UpperCAmelCase_							(    self				:str						,	lowercase_				:Any						,	lowercase_				:List[str]						,	lowercase_				:Any						,	lowercase_				:int						,	lowercase_				:List[Any]						,	lowercase_				:Tuple							)->							Tuple:
     A__    =      self.num_labels
     A__    =      TFDistilBertForTokenClassification(lowercase_							)
     A__    =      {"input_ids": input_ids, "attention_mask": input_mask}
     A__    =      model(lowercase_							)
     self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.seq_length, self.num_labels)							)



 def 					UpperCAmelCase_							(    self				:Any							)->							Union[str, Any]:
     A__    =      self.prepare_config_and_inputs()
     ((A__),     (A__),     (A__),     (A__),     (A__),     (A__))    =      config_and_inputs
     A__    =      {"input_ids": input_ids, "attention_mask": input_mask}
     return config, inputs_dict






@require_tf
class        UpperCAmelCase     (	UpperCamelCase__      ,	UpperCamelCase__      ,	unittest.TestCase    ):
 __lowercase       		=    (
     (
         TFDistilBertModel,
         TFDistilBertForMaskedLM,
         TFDistilBertForQuestionAnswering,
         TFDistilBertForSequenceClassification,
         TFDistilBertForTokenClassification,
         TFDistilBertForMultipleChoice,
     )
     if is_tf_available()
     else None
 )
 __lowercase       		=    (
     {
         """feature-extraction""": TFDistilBertModel,
         """fill-mask""": TFDistilBertForMaskedLM,
         """question-answering""": TFDistilBertForQuestionAnswering,
         """text-classification""": TFDistilBertForSequenceClassification,
         """token-classification""": TFDistilBertForTokenClassification,
         """zero-shot""": TFDistilBertForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 __lowercase       		=    False
 __lowercase       		=    False
 def 					UpperCAmelCase_							(    self				:Optional[Any]							)->							List[Any]:
     A__    =      TFDistilBertModelTester(self							)
     A__    =      ConfigTester(self						,	config_class=lowercase_						,	dim=37							)
 def 					UpperCAmelCase_							(    self				:Tuple							)->							Tuple:
     self.config_tester.run_common_tests()
 def 					UpperCAmelCase_							(    self				:int							)->							Tuple:
     A__    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_model(*lowercase_							)
 def 					UpperCAmelCase_							(    self				:Optional[int]							)->							Optional[Any]:
     A__    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_masked_lm(*lowercase_							)
 def 					UpperCAmelCase_							(    self				:str							)->							str:
     A__    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_question_answering(*lowercase_							)
 def 					UpperCAmelCase_							(    self				:List[str]							)->							Dict:
     A__    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowercase_							)
 def 					UpperCAmelCase_							(    self				:List[str]							)->							Optional[int]:
     A__    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowercase_							)
 def 					UpperCAmelCase_							(    self				:str							)->							int:
     A__    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_token_classification(*lowercase_							)



 @slow
 def 					UpperCAmelCase_							(    self				:List[str]							)->							Dict:
     for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]							):
         A__    =      TFDistilBertModel.from_pretrained(lowercase_							)
         self.assertIsNotNone(lowercase_							)






@require_tf
class        UpperCAmelCase     (	unittest.TestCase    ):
 @slow
 def 					UpperCAmelCase_							(    self				:List[Any]							)->							Any:
     A__    =      TFDistilBertModel.from_pretrained("distilbert-base-uncased"							)
     A__    =      tf.constant([[0, 1, 2, 3, 4, 5]]							)
     A__    =      model(lowercase_							)[0]

     A__    =      [1, 6, 7_68]
     self.assertEqual(output.shape						,	lowercase_							)

     A__    =      tf.constant(
         [
             [
                 [0.1_9_2_6_1_8_8_5, -0.1_3_7_3_2_9_5_5, 0.4_1_1_9_7_9_9],
                 [0.2_2_1_5_0_1_5_6, -0.0_7_4_2_2_6_6_1, 0.3_9_0_3_7_2_0_4],
                 [0.2_2_7_5_6_0_1_8, -0.0_8_9_6_4_1_4, 0.3_7_0_1_4_6_7],
             ]
         ]							)
     tf.debugging.assert_near(output[:, :3, :3]						,	lowercase_						,	atol=1E-4							)



