



from __future__ import annotations

import unittest

from transformers import AlbertConfig, is_tf_available
from transformers.models.auto import get_values
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 import TF_MODEL_FOR_PRETRAINING_MAPPING
	from transformers.models.albert.modeling_tf_albert import (
	    TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
	    TFAlbertForMaskedLM,
	    TFAlbertForMultipleChoice,
	    TFAlbertForPreTraining,
	    TFAlbertForQuestionAnswering,
	    TFAlbertForSequenceClassification,
	    TFAlbertForTokenClassification,
	    TFAlbertModel,
	)





class  __lowerCAmelCase			:




				def __init__(							self						,    _snake_case						,    _snake_case=13						,    _snake_case=7						,    _snake_case=True						,    _snake_case=True						,    _snake_case=True						,    _snake_case=True						,    _snake_case=99						,    _snake_case=16						,    _snake_case=32						,    _snake_case=2						,    _snake_case=4						,    _snake_case=37						,    _snake_case="gelu"						,    _snake_case=0.1						,    _snake_case=0.1						,    _snake_case=512						,    _snake_case=16						,    _snake_case=2						,    _snake_case=0.02						,    _snake_case=3						,    _snake_case=4						,    _snake_case=None						,    ):



									"""simple docstring"""





									_lowerCAmelCase      =   parent
									_lowerCAmelCase      =   13
									_lowerCAmelCase      =   7
									_lowerCAmelCase      =   True
									_lowerCAmelCase      =   True
									_lowerCAmelCase      =   True
									_lowerCAmelCase      =   True
									_lowerCAmelCase      =   99
									_lowerCAmelCase      =   16
									_lowerCAmelCase      =   32
									_lowerCAmelCase      =   2
									_lowerCAmelCase      =   4
									_lowerCAmelCase      =   37
									_lowerCAmelCase      =   """gelu"""
									_lowerCAmelCase      =   0.1
									_lowerCAmelCase      =   0.1
									_lowerCAmelCase      =   512
									_lowerCAmelCase      =   16
									_lowerCAmelCase      =   2
									_lowerCAmelCase      =   0.02
									_lowerCAmelCase      =   3
									_lowerCAmelCase      =   4
									_lowerCAmelCase      =   None




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   ids_tensor([self.batch_size, self.seq_length]						,    self.vocab_size		)

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

									_lowerCAmelCase      =   None
									if self.use_token_type_ids:
														_lowerCAmelCase      =   ids_tensor([self.batch_size, self.seq_length]						,    self.type_vocab_size		)

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

									_lowerCAmelCase      =   AlbertConfig(
									    vocab_size=self.vocab_size						,    hidden_size=self.hidden_size						,    embedding_size=self.embedding_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						,    type_vocab_size=self.type_vocab_size						,    initializer_range=self.initializer_range						,    )

									return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




				def 							snake_case       (							self						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case		):



									"""simple docstring"""





									_lowerCAmelCase      =   TFAlbertModel(config=_snake_case		)
									# inputs = {'input_ids': input_ids,
									#           'attention_mask': input_mask,
									#           'token_type_ids': token_type_ids}
									# sequence_output, pooled_output = model(**inputs)
									_lowerCAmelCase      =   {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
									_lowerCAmelCase      =   model(_snake_case		)

									_lowerCAmelCase      =   [input_ids, input_mask]
									_lowerCAmelCase      =   model(_snake_case		)

									_lowerCAmelCase      =   model(_snake_case		)

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




				def 							snake_case       (							self						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.num_labels
									_lowerCAmelCase      =   TFAlbertForPreTraining(config=_snake_case		)
									_lowerCAmelCase      =   {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
									_lowerCAmelCase      =   model(_snake_case		)
									self.parent.assertEqual(result.prediction_logits.shape						,    (self.batch_size, self.seq_length, self.vocab_size)		)
									self.parent.assertEqual(result.sop_logits.shape						,    (self.batch_size, self.num_labels)		)




				def 							snake_case       (							self						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case		):



									"""simple docstring"""





									_lowerCAmelCase      =   TFAlbertForMaskedLM(config=_snake_case		)
									_lowerCAmelCase      =   {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
									_lowerCAmelCase      =   model(_snake_case		)
									self.parent.assertEqual(result.logits.shape						,    (self.batch_size, self.seq_length, self.vocab_size)		)




				def 							snake_case       (							self						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.num_labels
									_lowerCAmelCase      =   TFAlbertForSequenceClassification(config=_snake_case		)
									_lowerCAmelCase      =   {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
									_lowerCAmelCase      =   model(_snake_case		)
									self.parent.assertEqual(result.logits.shape						,    (self.batch_size, self.num_labels)		)




				def 							snake_case       (							self						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case		):



									"""simple docstring"""





									_lowerCAmelCase      =   TFAlbertForQuestionAnswering(config=_snake_case		)
									_lowerCAmelCase      =   {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
									_lowerCAmelCase      =   model(_snake_case		)
									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 							snake_case       (							self						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.num_choices
									_lowerCAmelCase      =   TFAlbertForMultipleChoice(config=_snake_case		)
									_lowerCAmelCase      =   tf.tile(tf.expand_dims(_snake_case						,    1		)						,    (1, self.num_choices, 1)		)
									_lowerCAmelCase      =   tf.tile(tf.expand_dims(_snake_case						,    1		)						,    (1, self.num_choices, 1)		)
									_lowerCAmelCase      =   tf.tile(tf.expand_dims(_snake_case						,    1		)						,    (1, self.num_choices, 1)		)
									_lowerCAmelCase      =   {
									    """input_ids""": multiple_choice_inputs_ids,
									    """attention_mask""": multiple_choice_input_mask,
									    """token_type_ids""": multiple_choice_token_type_ids,
									}
									_lowerCAmelCase      =   model(_snake_case		)
									self.parent.assertListEqual(list(result["""logits"""].shape		)						,    [self.batch_size, self.num_choices]		)




				def 							snake_case       (							self						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case						,    _snake_case		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.num_labels
									_lowerCAmelCase      =   TFAlbertForTokenClassification(config=_snake_case		)
									_lowerCAmelCase      =   {
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """token_type_ids""": token_type_ids,
									}
									_lowerCAmelCase      =   model(_snake_case		)
									self.parent.assertListEqual(list(result["""logits"""].shape		)						,    [self.batch_size, self.seq_length, self.num_labels]		)




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.prepare_config_and_inputs()
									(
									    (
									    _lowerCAmelCase
									)     ,     (
									    _lowerCAmelCase
									)     ,     (
									    _lowerCAmelCase
									)     ,     (
									    _lowerCAmelCase
									)     ,     (
									    _lowerCAmelCase
									)     ,     (
									    _lowerCAmelCase
									)     ,     (
									    _lowerCAmelCase
									)     ,     
									)      =   config_and_inputs
									_lowerCAmelCase      =   {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
									return config, inputs_dict





@require_tf
class  __lowerCAmelCase			(					lowerCamelCase__	,						lowerCamelCase__	,						unittest.TestCase	):
				__lowerCamelCase			=    (
				    (
				        TFAlbertModel,
				        TFAlbertForPreTraining,
				        TFAlbertForMaskedLM,
				        TFAlbertForSequenceClassification,
				        TFAlbertForQuestionAnswering,
				        TFAlbertForTokenClassification,
				        TFAlbertForMultipleChoice,
				    )
				    if is_tf_available()
				    else ()
				)
				__lowerCamelCase			=    (
				    {
				        '''feature-extraction''': TFAlbertModel,
				        '''fill-mask''': TFAlbertForMaskedLM,
				        '''question-answering''': TFAlbertForQuestionAnswering,
				        '''text-classification''': TFAlbertForSequenceClassification,
				        '''token-classification''': TFAlbertForTokenClassification,
				        '''zero-shot''': TFAlbertForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				__lowerCamelCase			=    False
				__lowerCamelCase			=    False




				def 							snake_case       (							self						,    _snake_case						,    _snake_case						,    _snake_case=False		):



									"""simple docstring"""





									_lowerCAmelCase      =   super()._prepare_for_class(_snake_case						,    _snake_case						,    return_labels=_snake_case		)

									if return_labels:
														if model_class in get_values(_snake_case		):
																			_lowerCAmelCase      =   tf.zeros(self.model_tester.batch_size						,    dtype=tf.intaa		)

									return inputs_dict




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   TFAlbertModelTester(self		)
									_lowerCAmelCase      =   ConfigTester(self						,    config_class=_snake_case						,    hidden_size=37		)




				def 							snake_case       (							self		):



									"""simple docstring"""





									self.config_tester.run_common_tests()




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_albert_model(*_snake_case		)




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_albert_for_pretraining(*_snake_case		)




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_albert_for_masked_lm(*_snake_case		)




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_albert_for_multiple_choice(*_snake_case		)




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_albert_for_sequence_classification(*_snake_case		)




				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_albert_for_question_answering(*_snake_case		)




				@slow
				def 							snake_case       (							self		):



									"""simple docstring"""





									for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														_lowerCAmelCase      =   TFAlbertModel.from_pretrained(_snake_case		)
														self.assertIsNotNone(_snake_case		)





@require_tf
class  __lowerCAmelCase			(					unittest.TestCase	):




				@slow
				def 							snake_case       (							self		):



									"""simple docstring"""





									_lowerCAmelCase      =   TFAlbertForPreTraining.from_pretrained("""albert-base-v2"""		)
									_lowerCAmelCase      =   tf.constant([[0, 1, 2, 3, 4, 5]]		)
									_lowerCAmelCase      =   model(_snake_case		)[0]

									_lowerCAmelCase      =   [1, 6, 30000]
									self.assertEqual(output.shape						,    _snake_case		)

									_lowerCAmelCase      =   tf.constant(
									    [
									        [
									            [4.59_5668, 0.7446_2754, -1.81_8147],
									            [4.595_4347, 0.745_4184, -1.818_8258],
									            [4.595_4905, 0.744_8235, -1.818_2316],
									        ]
									    ]		)
									tf.debugging.assert_near(output[:, :3, :3]						,    _snake_case						,    atol=1e-4		)








# Lint as: python3

import hashlib
from typing import Union


def 		_UpperCAmelCase    (				snake_case      ):





					"""simple docstring"""

					if isinstance(snake_case     ,						snake_case      ):
										# Data already in bytes, returns as it as
										return hash_data
					elif isinstance(snake_case     ,						snake_case      ):
										# We keep the data as it as for it ot be later encoded to UTF-8
										# However replace `\\` with `/` for Windows compatibility
										_lowerCAmelCase      =   hash_data.replace("""\\"""     ,						"""/"""      )
					elif isinstance(snake_case     ,						snake_case      ):
										_lowerCAmelCase      =   str(snake_case      )
					else:
										# If data is not of the required type, raise error
										raise InvalidKeyError(snake_case      )

					return hash_data.encode("""utf-8"""      )





class  __lowerCAmelCase			(					lowerCamelCase__	):




				def __init__(							self						,    _snake_case		):



									"""simple docstring"""





									_lowerCAmelCase      =   """\nFAILURE TO GENERATE DATASET: Invalid key type detected"""
									_lowerCAmelCase      =   F'\nFound Key {hash_data} of type {type(_snake_case		)}'
									_lowerCAmelCase      =   """\nKeys should be either str, int or bytes type"""
									super().__init__(F'{self.prefix}{self.err_msg}{self.suffix}'		)





class  __lowerCAmelCase			(					lowerCamelCase__	):




				def __init__(							self						,    _snake_case						,    _snake_case						,    _snake_case=""		):



									"""simple docstring"""





									_lowerCAmelCase      =   key
									_lowerCAmelCase      =   duplicate_key_indices
									_lowerCAmelCase      =   fix_msg
									_lowerCAmelCase      =   """Found multiple examples generated with the same key"""
									if len(_snake_case		) <= 20:
														_lowerCAmelCase      =   F'\nThe examples at index {", ".join(_snake_case		)} have the key {key}'
									else:
														_lowerCAmelCase      =   F'\nThe examples at index {", ".join(duplicate_key_indices[:20]		)}... ({len(_snake_case		) - 20} more) have the key {key}'
									_lowerCAmelCase      =   """\n""" + fix_msg if fix_msg else """"""
									super().__init__(F'{self.prefix}{self.err_msg}{self.suffix}'		)





class  __lowerCAmelCase			:




				def __init__(							self						,    _snake_case		):



									"""simple docstring"""





									_lowerCAmelCase      =   hashlib.mda(_as_bytes(_snake_case		)		)




				def 							snake_case       (							self						,    _snake_case		):



									"""simple docstring"""





									_lowerCAmelCase      =   self._split_mda.copy()
									_lowerCAmelCase      =   _as_bytes(_snake_case		)
									mda.update(_snake_case		)
									# Convert to integer with hexadecimal conversion
									return int(mda.hexdigest()						,    16		)



