






"""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   SCREAMING_SNAKE_CASE	:





	"""simple docstring"""




	def __init__(self  ,  lowerCAmelCase_  ,  ):
		A_		:    Optional[Any]						     = parent
		A_		:    Dict						     = 13
		A_		:    List[Any]						     = 7
		A_		:    Union[str, Any]						     = True
		A_		:    int						     = True
		A_		:    str						     = False
		A_		:    List[Any]						     = True
		A_		:    Tuple						     = 99
		A_		:    Any						     = 32
		A_		:    Tuple						     = 2
		A_		:    Dict						     = 4
		A_		:    Any						     = 37
		A_		:    List[str]						     = """gelu"""
		A_		:    int						     = 0.1
		A_		:    List[Any]						     = 0.1
		A_		:    Dict						     = 512
		A_		:    Optional[int]						     = 16
		A_		:    Union[str, Any]						     = 2
		A_		:    str						     = 0.02
		A_		:    Dict						     = 3
		A_		:    Any						     = 4
		A_		:    Dict						     = None





	def  lowerCamelCase(self				):
		A_		:    int						     = ids_tensor([self.batch_size, self.seq_length]  ,  self.vocab_size				)

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

		A_		:    Any						     = None
		A_		:    Tuple						     = None
		A_		:    List[Any]						     = None
		if self.use_labels:
			A_		:    str						     = ids_tensor([self.batch_size]  ,  self.type_sequence_label_size				)
			A_		:    Dict						     = ids_tensor([self.batch_size, self.seq_length]  ,  self.num_labels				)
			A_		:    Tuple						     = ids_tensor([self.batch_size]  ,  self.num_choices				)

		A_		:    List[Any]						     = 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  lowerCamelCase(self  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_				):
		A_		:    Any						     = TFDistilBertModel(config=lowerCAmelCase_				)
		A_		:    Optional[Any]						     = {"""input_ids""": input_ids, """attention_mask""": input_mask}

		A_		:    Any						     = model(lowerCAmelCase_				)

		A_		:    Dict						     = [input_ids, input_mask]

		A_		:    Tuple						     = model(lowerCAmelCase_				)

		self.parent.assertEqual(result.last_hidden_state.shape  ,  (self.batch_size, self.seq_length, self.hidden_size)				)





	def  lowerCamelCase(self  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_				):
		A_		:    Dict						     = TFDistilBertForMaskedLM(config=lowerCAmelCase_				)
		A_		:    Optional[Any]						     = {"""input_ids""": input_ids, """attention_mask""": input_mask}
		A_		:    Tuple						     = model(lowerCAmelCase_				)
		self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.vocab_size)				)





	def  lowerCamelCase(self  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_				):
		A_		:    Any						     = TFDistilBertForQuestionAnswering(config=lowerCAmelCase_				)
		A_		:    int						     = {
		    """input_ids""": input_ids,
		    """attention_mask""": input_mask,
		}
		A_		:    str						     = model(lowerCAmelCase_				)
		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  lowerCamelCase(self  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_				):
		A_		:    Optional[Any]						     = self.num_labels
		A_		:    Optional[int]						     = TFDistilBertForSequenceClassification(lowerCAmelCase_				)
		A_		:    int						     = {"""input_ids""": input_ids, """attention_mask""": input_mask}
		A_		:    List[Any]						     = model(lowerCAmelCase_				)
		self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_labels)				)





	def  lowerCamelCase(self  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_				):
		A_		:    int						     = self.num_choices
		A_		:    str						     = TFDistilBertForMultipleChoice(lowerCAmelCase_				)
		A_		:    str						     = tf.tile(tf.expand_dims(lowerCAmelCase_  ,  1				)  ,  (1, self.num_choices, 1)				)
		A_		:    int						     = tf.tile(tf.expand_dims(lowerCAmelCase_  ,  1				)  ,  (1, self.num_choices, 1)				)
		A_		:    Dict						     = {
		    """input_ids""": multiple_choice_inputs_ids,
		    """attention_mask""": multiple_choice_input_mask,
		}
		A_		:    List[str]						     = model(lowerCAmelCase_				)
		self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_choices)				)





	def  lowerCamelCase(self  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_  ,  lowerCAmelCase_				):
		A_		:    List[Any]						     = self.num_labels
		A_		:    List[Any]						     = TFDistilBertForTokenClassification(lowerCAmelCase_				)
		A_		:    Optional[Any]						     = {"""input_ids""": input_ids, """attention_mask""": input_mask}
		A_		:    str						     = model(lowerCAmelCase_				)
		self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.num_labels)				)





	def  lowerCamelCase(self				):
		A_		:    Tuple						     = self.prepare_config_and_inputs()
		((A_)     ,  (A_)     ,  (A_)     ,  (A_)     ,  (A_)     ,  (A_))		:    Tuple						     = config_and_inputs
		A_		:    Optional[int]						     = {"""input_ids""": input_ids, """attention_mask""": input_mask}
		return config, inputs_dict






@require_tf
class   SCREAMING_SNAKE_CASE	(					_SCREAMING_SNAKE_CASE							,							_SCREAMING_SNAKE_CASE							,							unittest.TestCase    ):





	"""simple docstring"""




	_A							:					List[Any] 			=					(
	    (
	        TFDistilBertModel,
	        TFDistilBertForMaskedLM,
	        TFDistilBertForQuestionAnswering,
	        TFDistilBertForSequenceClassification,
	        TFDistilBertForTokenClassification,
	        TFDistilBertForMultipleChoice,
	    )
	    if is_tf_available()
	    else None
	)
	_A							:					Dict 			=					(
	    {
	        """feature-extraction""": TFDistilBertModel,
	        """fill-mask""": TFDistilBertForMaskedLM,
	        """question-answering""": TFDistilBertForQuestionAnswering,
	        """text-classification""": TFDistilBertForSequenceClassification,
	        """token-classification""": TFDistilBertForTokenClassification,
	        """zero-shot""": TFDistilBertForSequenceClassification,
	    }
	    if is_tf_available()
	    else {}
	)
	_A							:					Union[str, Any] 			=					False
	_A							:					Any 			=					False





	def  lowerCamelCase(self				):
		A_		:    Tuple						     = TFDistilBertModelTester(self				)
		A_		:    List[Any]						     = ConfigTester(self  ,  config_class=lowerCAmelCase_  ,  dim=37				)





	def  lowerCamelCase(self				):
		self.config_tester.run_common_tests()





	def  lowerCamelCase(self				):
		A_		:    str						     = self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_distilbert_model(*lowerCAmelCase_				)





	def  lowerCamelCase(self				):
		A_		:    Optional[int]						     = self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_distilbert_for_masked_lm(*lowerCAmelCase_				)





	def  lowerCamelCase(self				):
		A_		:    Dict						     = self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_distilbert_for_question_answering(*lowerCAmelCase_				)





	def  lowerCamelCase(self				):
		A_		:    Optional[Any]						     = self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowerCAmelCase_				)





	def  lowerCamelCase(self				):
		A_		:    Any						     = self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowerCAmelCase_				)





	def  lowerCamelCase(self				):
		A_		:    str						     = self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_distilbert_for_token_classification(*lowerCAmelCase_				)





	@slow
	def  lowerCamelCase(self				):
		for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]				):
			A_		:    Dict						     = TFDistilBertModel.from_pretrained(lowerCAmelCase_				)
			self.assertIsNotNone(lowerCAmelCase_				)






@require_tf
class   SCREAMING_SNAKE_CASE	(					unittest.TestCase    ):





	"""simple docstring"""




	@slow
	def  lowerCamelCase(self				):
		A_		:    List[str]						     = TFDistilBertModel.from_pretrained("""distilbert-base-uncased"""				)
		A_		:    Any						     = tf.constant([[0, 1, 2, 3, 4, 5]]				)
		A_		:    Optional[int]						     = model(lowerCAmelCase_				)[0]

		A_		:    str						     = [1, 6, 768]
		self.assertEqual(output.shape  ,  lowerCAmelCase_				)

		A_		:    Optional[Any]						     = tf.constant(
		    [
		        [
		            [0.19261885, -0.13732955, 0.4119799],
		            [0.22150156, -0.07422661, 0.39037204],
		            [0.22756018, -0.0896414, 0.3701467],
		        ]
		    ]				)
		tf.debugging.assert_near(output[:, :3, :3]  ,  lowerCAmelCase_  ,  atol=1e-4				)









"""simple docstring"""
from typing import Dict

from .base import GenericTensor, Pipeline



class   SCREAMING_SNAKE_CASE	(					_SCREAMING_SNAKE_CASE    ):





	"""simple docstring"""




	def  lowerCamelCase(self  ,  lowerCAmelCase_=None  ,  lowerCAmelCase_=None  ,  lowerCAmelCase_=None  ,  **lowerCAmelCase_				):
		if tokenize_kwargs is None:
			A_		:    Dict						     = {}

		if truncation is not None:
			if "truncation" in tokenize_kwargs:
				raise ValueError(
				    """truncation parameter defined twice (given as keyword argument as well as in tokenize_kwargs)"""				)
			A_		:    Optional[Any]						     = truncation

		A_		:    Optional[int]						     = tokenize_kwargs

		A_		:    Optional[Any]						     = {}
		if return_tensors is not None:
			A_		:    Optional[Any]						     = return_tensors

		return preprocess_params, {}, postprocess_params





	def  lowerCamelCase(self  ,  lowerCAmelCase_  ,  **lowerCAmelCase_				):
		A_		:    List[Any]						     = self.framework
		A_		:    Optional[int]						     = self.tokenizer(lowerCAmelCase_  ,  return_tensors=lowerCAmelCase_  ,  **lowerCAmelCase_				)
		return model_inputs





	def  lowerCamelCase(self  ,  lowerCAmelCase_				):
		A_		:    List[Any]						     = self.model(**lowerCAmelCase_				)
		return model_outputs





	def  lowerCamelCase(self  ,  lowerCAmelCase_  ,  lowerCAmelCase_=False				):
		# [0] is the first available tensor, logits or last_hidden_state.
		if return_tensors:
			return model_outputs[0]
		if self.framework == "pt":
			return model_outputs[0].tolist()
		elif self.framework == "tf":
			return model_outputs[0].numpy().tolist()





	def __call__(self  ,  *lowerCAmelCase_  ,  **lowerCAmelCase_				):
		return super().__call__(*lowerCAmelCase_  ,  **lowerCAmelCase_				)

