

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 	lowerCamelCase      :







					'''simple docstring'''







					def __init__(			self    ,	lowerCAmelCase    ,	):
							UpperCAmelCase_						=				parent
							UpperCAmelCase_						=				13
							UpperCAmelCase_						=				7
							UpperCAmelCase_						=				True
							UpperCAmelCase_						=				True
							UpperCAmelCase_						=				False
							UpperCAmelCase_						=				True
							UpperCAmelCase_						=				99
							UpperCAmelCase_						=				32
							UpperCAmelCase_						=				2
							UpperCAmelCase_						=				4
							UpperCAmelCase_						=				37
							UpperCAmelCase_						=				'''gelu'''
							UpperCAmelCase_						=				0.1
							UpperCAmelCase_						=				0.1
							UpperCAmelCase_						=				512
							UpperCAmelCase_						=				16
							UpperCAmelCase_						=				2
							UpperCAmelCase_						=				0.02
							UpperCAmelCase_						=				3
							UpperCAmelCase_						=				4
							UpperCAmelCase_						=				None



					def    A__		(			self     ):
							UpperCAmelCase_						=				ids_tensor([self.batch_size, self.seq_length]    ,	self.vocab_size     )

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

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

							UpperCAmelCase_						=				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    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
							UpperCAmelCase_						=				TFDistilBertModel(config=UpperCAmelCase__     )
							UpperCAmelCase_						=				{'''input_ids''': input_ids, '''attention_mask''': input_mask}

							UpperCAmelCase_						=				model(UpperCAmelCase__     )

							UpperCAmelCase_						=				[input_ids, input_mask]

							UpperCAmelCase_						=				model(UpperCAmelCase__     )

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



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
							UpperCAmelCase_						=				TFDistilBertForMaskedLM(config=UpperCAmelCase__     )
							UpperCAmelCase_						=				{'''input_ids''': input_ids, '''attention_mask''': input_mask}
							UpperCAmelCase_						=				model(UpperCAmelCase__     )
							self.parent.assertEqual(result.logits.shape    ,	(self.batch_size, self.seq_length, self.vocab_size)     )



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
							UpperCAmelCase_						=				TFDistilBertForQuestionAnswering(config=UpperCAmelCase__     )
							UpperCAmelCase_						=				{
							    '''input_ids''': input_ids,
							    '''attention_mask''': input_mask,
							}
							UpperCAmelCase_						=				model(UpperCAmelCase__     )
							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    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
							UpperCAmelCase_						=				self.num_labels
							UpperCAmelCase_						=				TFDistilBertForSequenceClassification(UpperCAmelCase__     )
							UpperCAmelCase_						=				{'''input_ids''': input_ids, '''attention_mask''': input_mask}
							UpperCAmelCase_						=				model(UpperCAmelCase__     )
							self.parent.assertEqual(result.logits.shape    ,	(self.batch_size, self.num_labels)     )



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
							UpperCAmelCase_						=				self.num_choices
							UpperCAmelCase_						=				TFDistilBertForMultipleChoice(UpperCAmelCase__     )
							UpperCAmelCase_						=				tf.tile(tf.expand_dims(UpperCAmelCase__    ,	1     )    ,	(1, self.num_choices, 1)     )
							UpperCAmelCase_						=				tf.tile(tf.expand_dims(UpperCAmelCase__    ,	1     )    ,	(1, self.num_choices, 1)     )
							UpperCAmelCase_						=				{
							    '''input_ids''': multiple_choice_inputs_ids,
							    '''attention_mask''': multiple_choice_input_mask,
							}
							UpperCAmelCase_						=				model(UpperCAmelCase__     )
							self.parent.assertEqual(result.logits.shape    ,	(self.batch_size, self.num_choices)     )



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
							UpperCAmelCase_						=				self.num_labels
							UpperCAmelCase_						=				TFDistilBertForTokenClassification(UpperCAmelCase__     )
							UpperCAmelCase_						=				{'''input_ids''': input_ids, '''attention_mask''': input_mask}
							UpperCAmelCase_						=				model(UpperCAmelCase__     )
							self.parent.assertEqual(result.logits.shape    ,	(self.batch_size, self.seq_length, self.num_labels)     )



					def    A__		(			self     ):
							UpperCAmelCase_						=				self.prepare_config_and_inputs()
							(UpperCAmelCase_)						=				config_and_inputs
							UpperCAmelCase_						=				{'''input_ids''': input_ids, '''attention_mask''': input_mask}
							return config, inputs_dict





@require_tf
class 	lowerCamelCase      (						lowercase__,     lowercase__,     unittest.TestCase  ):







					'''simple docstring'''







					lowerCAmelCase_				:    str									=  (
					    (
					        TFDistilBertModel,
					        TFDistilBertForMaskedLM,
					        TFDistilBertForQuestionAnswering,
					        TFDistilBertForSequenceClassification,
					        TFDistilBertForTokenClassification,
					        TFDistilBertForMultipleChoice,
					    )
					    if is_tf_available()
					    else None
					)
					lowerCAmelCase_				:    int									=  (
					    {
					        'feature-extraction': TFDistilBertModel,
					        'fill-mask': TFDistilBertForMaskedLM,
					        'question-answering': TFDistilBertForQuestionAnswering,
					        'text-classification': TFDistilBertForSequenceClassification,
					        'token-classification': TFDistilBertForTokenClassification,
					        'zero-shot': TFDistilBertForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					lowerCAmelCase_				:    Optional[Any]									=  False
					lowerCAmelCase_				:    int									=  False



					def    A__		(			self     ):
							UpperCAmelCase_						=				TFDistilBertModelTester(self     )
							UpperCAmelCase_						=				ConfigTester(self    ,	config_class=UpperCAmelCase__    ,	dim=37     )



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



					def    A__		(			self     ):
							UpperCAmelCase_						=				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_model(*UpperCAmelCase__     )



					def    A__		(			self     ):
							UpperCAmelCase_						=				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_masked_lm(*UpperCAmelCase__     )



					def    A__		(			self     ):
							UpperCAmelCase_						=				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_question_answering(*UpperCAmelCase__     )



					def    A__		(			self     ):
							UpperCAmelCase_						=				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_sequence_classification(*UpperCAmelCase__     )



					def    A__		(			self     ):
							UpperCAmelCase_						=				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_multiple_choice(*UpperCAmelCase__     )



					def    A__		(			self     ):
							UpperCAmelCase_						=				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_token_classification(*UpperCAmelCase__     )



					@slow
					def    A__		(			self     ):
							for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]     ):
									UpperCAmelCase_						=				TFDistilBertModel.from_pretrained(UpperCAmelCase__     )
									self.assertIsNotNone(UpperCAmelCase__     )





@require_tf
class 	lowerCamelCase      (						unittest.TestCase  ):







					'''simple docstring'''







					@slow
					def    A__		(			self     ):
							UpperCAmelCase_						=				TFDistilBertModel.from_pretrained("distilbert-base-uncased"     )
							UpperCAmelCase_						=				tf.constant([[0, 1, 2, 3, 4, 5]]     )
							UpperCAmelCase_						=				model(UpperCAmelCase__     )[0]

							UpperCAmelCase_						=				[1, 6, 768]
							self.assertEqual(output.shape    ,	UpperCAmelCase__     )

							UpperCAmelCase_						=				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]    ,	UpperCAmelCase__    ,	atol=1e-4     )






import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
			from transformers import AutoModelForSeqaSeqLM, AutoTokenizer



@require_torch
@require_sentencepiece
@require_tokenizers
class 	lowerCamelCase      (						unittest.TestCase  ):







					'''simple docstring'''







					@slow
					def    A__		(			self     ):
							UpperCAmelCase_						=				AutoModelForSeqaSeqLM.from_pretrained("google/mt5-small"    ,	return_dict=lowerCAmelCase     ).to(lowerCAmelCase     )
							UpperCAmelCase_						=				AutoTokenizer.from_pretrained("google/mt5-small"     )

							UpperCAmelCase_						=				tokenizer("Hello there"    ,	return_tensors="pt"     ).input_ids
							UpperCAmelCase_						=				tokenizer("Hi I am"    ,	return_tensors="pt"     ).input_ids

							UpperCAmelCase_						=				model(input_ids.to(lowerCAmelCase     )    ,	labels=labels.to(lowerCAmelCase     )     ).loss
							UpperCAmelCase_						=				-(labels.shape[-1] * loss.item())

							UpperCAmelCase_						=				-84.9127
							self.assertTrue(abs(mtf_score - EXPECTED_SCORE     ) < 1e-4     )

