



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



						def __init__(self						,     _lowercase						,     ):





													'''simple docstring'''

													__a		:		List[Any]									=       parent
													__a		:		Optional[int]									=       13
													__a		:		Optional[int]									=       7
													__a		:		Any									=       True
													__a		:		Optional[int]									=       True
													__a		:		Any									=       False
													__a		:		str									=       True
													__a		:		Optional[Any]									=       99
													__a		:		Dict									=       32
													__a		:		List[Any]									=       2
													__a		:		int									=       4
													__a		:		List[str]									=       37
													__a		:		Dict									=       """gelu"""
													__a		:		Any									=       0.1
													__a		:		Tuple									=       0.1
													__a		:		Tuple									=       512
													__a		:		Optional[int]									=       16
													__a		:		Any									=       2
													__a		:		Optional[int]									=       0.02
													__a		:		str									=       3
													__a		:		Optional[int]									=       4
													__a		:		List[str]									=       None



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Any									=       ids_tensor([self.batch_size, self.seq_length]						,     self.vocab_size							)

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

													__a		:		str									=       None
													__a		:		List[Any]									=       None
													__a		:		Union[str, Any]									=       None
													if self.use_labels:
																				__a		:		Optional[int]									=       ids_tensor([self.batch_size]						,     self.type_sequence_label_size							)
																				__a		:		Optional[int]									=       ids_tensor([self.batch_size, self.seq_length]						,     self.num_labels							)
																				__a		:		Optional[int]									=       ids_tensor([self.batch_size]						,     self.num_choices							)

													__a		:		List[str]									=       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						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							):





													'''simple docstring'''

													__a		:		int									=       TFDistilBertModel(config=_lowercase							)
													__a		:		Optional[Any]									=       {"""input_ids""": input_ids, """attention_mask""": input_mask}

													__a		:		str									=       model(_lowercase							)

													__a		:		Optional[Any]									=       [input_ids, input_mask]

													__a		:		Any									=       model(_lowercase							)

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



						def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							):





													'''simple docstring'''

													__a		:		Dict									=       TFDistilBertForMaskedLM(config=_lowercase							)
													__a		:		List[str]									=       {"""input_ids""": input_ids, """attention_mask""": input_mask}
													__a		:		Any									=       model(_lowercase							)
													self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.seq_length, self.vocab_size)							)



						def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							):





													'''simple docstring'''

													__a		:		Optional[int]									=       TFDistilBertForQuestionAnswering(config=_lowercase							)
													__a		:		str									=       {
													    """input_ids""": input_ids,
													    """attention_mask""": input_mask,
													}
													__a		:		Tuple									=       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     lowerCAmelCase__(self						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							):





													'''simple docstring'''

													__a		:		Optional[int]									=       self.num_labels
													__a		:		Union[str, Any]									=       TFDistilBertForSequenceClassification(_lowercase							)
													__a		:		str									=       {"""input_ids""": input_ids, """attention_mask""": input_mask}
													__a		:		Tuple									=       model(_lowercase							)
													self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.num_labels)							)



						def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							):





													'''simple docstring'''

													__a		:		Optional[int]									=       self.num_choices
													__a		:		List[Any]									=       TFDistilBertForMultipleChoice(_lowercase							)
													__a		:		List[str]									=       tf.tile(tf.expand_dims(_lowercase						,     1							)						,     (1, self.num_choices, 1)							)
													__a		:		int									=       tf.tile(tf.expand_dims(_lowercase						,     1							)						,     (1, self.num_choices, 1)							)
													__a		:		int									=       {
													    """input_ids""": multiple_choice_inputs_ids,
													    """attention_mask""": multiple_choice_input_mask,
													}
													__a		:		Optional[int]									=       model(_lowercase							)
													self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.num_choices)							)



						def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     _lowercase							):





													'''simple docstring'''

													__a		:		List[str]									=       self.num_labels
													__a		:		Tuple									=       TFDistilBertForTokenClassification(_lowercase							)
													__a		:		Any									=       {"""input_ids""": input_ids, """attention_mask""": input_mask}
													__a		:		List[str]									=       model(_lowercase							)
													self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.seq_length, self.num_labels)							)





						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Any									=       self.prepare_config_and_inputs()
													((__a)							,     (__a)							,     (__a)							,     (__a)							,     (__a)							,     (__a))		:		List[str]									=       config_and_inputs
													__a		:		Union[str, Any]									=       {"""input_ids""": input_ids, """attention_mask""": input_mask}
													return config, inputs_dict







@require_tf
class SCREAMING_SNAKE_CASE__					(		__snake_case     ,					__snake_case     ,					unittest.TestCase      ):
						_lowerCAmelCase							=					(
						    (
						        TFDistilBertModel,
						        TFDistilBertForMaskedLM,
						        TFDistilBertForQuestionAnswering,
						        TFDistilBertForSequenceClassification,
						        TFDistilBertForTokenClassification,
						        TFDistilBertForMultipleChoice,
						    )
						    if is_tf_available()
						    else None
						)
						_lowerCAmelCase							=					(
						    {
						        "feature-extraction": TFDistilBertModel,
						        "fill-mask": TFDistilBertForMaskedLM,
						        "question-answering": TFDistilBertForQuestionAnswering,
						        "text-classification": TFDistilBertForSequenceClassification,
						        "token-classification": TFDistilBertForTokenClassification,
						        "zero-shot": TFDistilBertForSequenceClassification,
						    }
						    if is_tf_available()
						    else {}
						)
						_lowerCAmelCase							=					False
						_lowerCAmelCase							=					False



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		List[str]									=       TFDistilBertModelTester(self							)
													__a		:		Optional[int]									=       ConfigTester(self						,     config_class=_lowercase						,     dim=37							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													self.config_tester.run_common_tests()



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Optional[int]									=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_distilbert_model(*_lowercase							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		List[Any]									=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_distilbert_for_masked_lm(*_lowercase							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		int									=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_distilbert_for_question_answering(*_lowercase							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Dict									=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_distilbert_for_sequence_classification(*_lowercase							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Union[str, Any]									=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_distilbert_for_multiple_choice(*_lowercase							)



						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Optional[Any]									=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_distilbert_for_token_classification(*_lowercase							)





						@slow
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]							):
																				__a		:		Any									=       TFDistilBertModel.from_pretrained(_lowercase							)
																				self.assertIsNotNone(_lowercase							)







@require_tf
class SCREAMING_SNAKE_CASE__					(		unittest.TestCase      ):



						@slow
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													__a		:		Union[str, Any]									=       TFDistilBertModel.from_pretrained("""distilbert-base-uncased"""							)
													__a		:		Union[str, Any]									=       tf.constant([[0, 1, 2, 3, 4, 5]]							)
													__a		:		Optional[Any]									=       model(_lowercase							)[0]

													__a		:		List[Any]									=       [1, 6, 768]
													self.assertEqual(output.shape						,     _lowercase							)

													__a		:		int									=       tf.constant(
													    [
													        [
													            [0.1926_1885, -0.1373_2955, 0.411_9799],
													            [0.2215_0156, -0.0742_2661, 0.3903_7204],
													            [0.2275_6018, -0.089_6414, 0.370_1467],
													        ]
													    ]							)
													tf.debugging.assert_near(output[:, :3, :3]						,     _lowercase						,     atol=1e-4							)






"""simple docstring"""

from typing import Callable, Dict, Optional, Tuple

import torch
from torch import nn
from torch.distributions import (
    AffineTransform,
    Distribution,
    Independent,
    NegativeBinomial,
    Normal,
    StudentT,
    TransformedDistribution,
)




class SCREAMING_SNAKE_CASE__					(		__snake_case      ):



						def __init__(self						,     _lowercase						,     _lowercase=None						,     _lowercase=None						,     _lowercase=0							):





													'''simple docstring'''

													__a		:		Any									=       1.0 if scale is None else scale
													__a		:		str									=       0.0 if loc is None else loc

													super().__init__(_lowercase						,     [AffineTransform(loc=self.loc						,     scale=self.scale						,     event_dim=_lowercase							)]							)



						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													return self.base_dist.mean * self.scale + self.loc



						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													return self.base_dist.variance * self.scale**2





						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													return self.variance.sqrt()







class SCREAMING_SNAKE_CASE__					(		nn.Module      ):



						def __init__(self						,     _lowercase						,     _lowercase						,     _lowercase						,     **_lowercase							):





													'''simple docstring'''

													super().__init__(**_lowercase							)
													__a		:		str									=       args_dim
													__a		:		List[Any]									=       nn.ModuleList([nn.Linear(_lowercase						,     _lowercase							) for dim in args_dim.values()]							)
													__a		:		Dict									=       domain_map





						def     lowerCAmelCase__(self						,     _lowercase							):





													'''simple docstring'''

													__a		:		List[Any]									=       [proj(_lowercase							) for proj in self.proj]

													return self.domain_map(*_lowercase							)







class SCREAMING_SNAKE_CASE__					(		nn.Module      ):



						def __init__(self						,     _lowercase							):





													'''simple docstring'''

													super().__init__()
													__a		:		Optional[int]									=       function





						def     lowerCAmelCase__(self						,     _lowercase						,     *_lowercase							):





													'''simple docstring'''

													return self.function(_lowercase						,     *_lowercase							)







class SCREAMING_SNAKE_CASE__					:
						_lowerCAmelCase							=					42
						_lowerCAmelCase							=					42
						_lowerCAmelCase							=					42



						def __init__(self						,     _lowercase = 1							):





													'''simple docstring'''

													__a		:		Optional[int]									=       dim
													__a		:		str									=       {k: dim * self.args_dim[k] for k in self.args_dim}



						def     lowerCAmelCase__(self						,     _lowercase							):





													'''simple docstring'''

													if self.dim == 1:
																				return self.distribution_class(*_lowercase							)
													else:
																				return Independent(self.distribution_class(*_lowercase							)						,     1							)



						def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase = None						,     _lowercase = None						,     ):





													'''simple docstring'''

													__a		:		Tuple									=       self._base_distribution(_lowercase							)
													if loc is None and scale is None:
																				return distr
													else:
																				return AffineTransformed(_lowercase						,     loc=_lowercase						,     scale=_lowercase						,     event_dim=self.event_dim							)



						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													return () if self.dim == 1 else (self.dim,)



						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													return len(self.event_shape							)



						@property
						def     lowerCAmelCase__(self							):





													'''simple docstring'''

													return 0.0



						def     lowerCAmelCase__(self						,     _lowercase							):





													'''simple docstring'''

													return ParameterProjection(
													    in_features=_lowercase						,     args_dim=self.args_dim						,     domain_map=LambdaLayer(self.domain_map							)						,     )



						def     lowerCAmelCase__(self						,     *_lowercase							):





													'''simple docstring'''

													raise NotImplementedError()





						@staticmethod
						def     lowerCAmelCase__(_lowercase							):





													'''simple docstring'''

													return (x + torch.sqrt(torch.square(_lowercase							) + 4.0							)) / 2.0







class SCREAMING_SNAKE_CASE__					(		__snake_case      ):
						_lowerCAmelCase							=					{"df": 1, "loc": 1, "scale": 1}
						_lowerCAmelCase							=					StudentT



						@classmethod
						def     lowerCAmelCase__(cls						,     _lowercase						,     _lowercase						,     _lowercase							):





													'''simple docstring'''

													__a		:		int									=       cls.squareplus(_lowercase							).clamp_min(torch.finfo(scale.dtype							).eps							)
													__a		:		Optional[Any]									=       2.0 + cls.squareplus(_lowercase							)
													return df.squeeze(-1							), loc.squeeze(-1							), scale.squeeze(-1							)







class SCREAMING_SNAKE_CASE__					(		__snake_case      ):
						_lowerCAmelCase							=					{"loc": 1, "scale": 1}
						_lowerCAmelCase							=					Normal



						@classmethod
						def     lowerCAmelCase__(cls						,     _lowercase						,     _lowercase							):





													'''simple docstring'''

													__a		:		str									=       cls.squareplus(_lowercase							).clamp_min(torch.finfo(scale.dtype							).eps							)
													return loc.squeeze(-1							), scale.squeeze(-1							)







class SCREAMING_SNAKE_CASE__					(		__snake_case      ):
						_lowerCAmelCase							=					{"total_count": 1, "logits": 1}
						_lowerCAmelCase							=					NegativeBinomial



						@classmethod
						def     lowerCAmelCase__(cls						,     _lowercase						,     _lowercase							):





													'''simple docstring'''

													__a		:		Union[str, Any]									=       cls.squareplus(_lowercase							)
													return total_count.squeeze(-1							), logits.squeeze(-1							)



						def     lowerCAmelCase__(self						,     _lowercase							):





													'''simple docstring'''

													__a							,     __a		:		Optional[Any]									=       distr_args
													if self.dim == 1:
																				return self.distribution_class(total_count=_lowercase						,     logits=_lowercase							)
													else:
																				return Independent(self.distribution_class(total_count=_lowercase						,     logits=_lowercase							)						,     1							)





						def     lowerCAmelCase__(self						,     _lowercase						,     _lowercase = None						,     _lowercase = None							):





													'''simple docstring'''

													__a							,     __a		:		List[Any]									=       distr_args

													if scale is not None:
																				# See scaling property of Gamma.
																				logits += scale.log()

													return self._base_distribution((total_count, logits)							)

