


import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
				import jax.numpy as jnp

				from transformers.models.albert.modeling_flax_albert import (
				    FlaxAlbertForMaskedLM,
				    FlaxAlbertForMultipleChoice,
				    FlaxAlbertForPreTraining,
				    FlaxAlbertForQuestionAnswering,
				    FlaxAlbertForSequenceClassification,
				    FlaxAlbertForTokenClassification,
				    FlaxAlbertModel,
				)




class 							lowerCamelCase   (  unittest.TestCase ):


	def __init__(    self				,     lowercase__				,     lowercase__=1_3				,     lowercase__=7				,     lowercase__=True				,     lowercase__=True				,     lowercase__=True				,     lowercase__=True				,     lowercase__=9_9				,     lowercase__=3_2				,     lowercase__=5				,     lowercase__=4				,     lowercase__=3_7				,     lowercase__="gelu"				,     lowercase__=0.1				,     lowercase__=0.1				,     lowercase__=5_1_2				,     lowercase__=1_6				,     lowercase__=2				,     lowercase__=0.0_2				,     lowercase__=4				,     ):
							__UpperCAmelCase			:							Optional[Any]	    =	parent
							__UpperCAmelCase			:							Optional[Any]	    =	batch_size
							__UpperCAmelCase			:							Optional[int]	    =	seq_length
							__UpperCAmelCase			:							Optional[Any]	    =	is_training
							__UpperCAmelCase			:							List[Any]	    =	use_attention_mask
							__UpperCAmelCase			:							Optional[int]	    =	use_token_type_ids
							__UpperCAmelCase			:							Optional[int]	    =	use_labels
							__UpperCAmelCase			:							List[Any]	    =	vocab_size
							__UpperCAmelCase			:							List[str]	    =	hidden_size
							__UpperCAmelCase			:							Union[str, Any]	    =	num_hidden_layers
							__UpperCAmelCase			:							Optional[int]	    =	num_attention_heads
							__UpperCAmelCase			:							List[Any]	    =	intermediate_size
							__UpperCAmelCase			:							str	    =	hidden_act
							__UpperCAmelCase			:							List[str]	    =	hidden_dropout_prob
							__UpperCAmelCase			:							Union[str, Any]	    =	attention_probs_dropout_prob
							__UpperCAmelCase			:							Optional[int]	    =	max_position_embeddings
							__UpperCAmelCase			:							Dict	    =	type_vocab_size
							__UpperCAmelCase			:							int	    =	type_sequence_label_size
							__UpperCAmelCase			:							Optional[Any]	    =	initializer_range
							__UpperCAmelCase			:							Dict	    =	num_choices


	def  A(    self):
							__UpperCAmelCase			:							Dict	    =	ids_tensor([self.batch_size, self.seq_length]				,     self.vocab_size)

							__UpperCAmelCase			:							Tuple	    =	None
							if self.use_attention_mask:
													__UpperCAmelCase			:							Any	    =	random_attention_mask([self.batch_size, self.seq_length])

							__UpperCAmelCase			:							Optional[int]	    =	None
							if self.use_token_type_ids:
													__UpperCAmelCase			:							Optional[Any]	    =	ids_tensor([self.batch_size, self.seq_length]				,     self.type_vocab_size)

							__UpperCAmelCase			:							Tuple	    =	AlbertConfig(
							    vocab_size=self.vocab_size				,     hidden_size=self.hidden_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				,     is_decoder=lowercase__				,     initializer_range=self.initializer_range				,     )

							return config, input_ids, token_type_ids, attention_mask


	def  A(    self):
							__UpperCAmelCase			:							Tuple	    =	self.prepare_config_and_inputs()
							__UpperCAmelCase    ,      __UpperCAmelCase    ,      __UpperCAmelCase    ,      __UpperCAmelCase			:							Optional[Any]	    =	config_and_inputs
							__UpperCAmelCase			:							int	    =	{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
							return config, inputs_dict






@require_flax
class 							lowerCamelCase   (  _UpperCamelCase							,						unittest.TestCase ):
	_lowerCAmelCase     :		Dict										=						(
	    (
	        FlaxAlbertModel,
	        FlaxAlbertForPreTraining,
	        FlaxAlbertForMaskedLM,
	        FlaxAlbertForMultipleChoice,
	        FlaxAlbertForQuestionAnswering,
	        FlaxAlbertForSequenceClassification,
	        FlaxAlbertForTokenClassification,
	        FlaxAlbertForQuestionAnswering,
	    )
	    if is_flax_available()
	    else ()
	)


	def  A(    self):
							__UpperCAmelCase			:							Any	    =	FlaxAlbertModelTester(self)


	@slow
	def  A(    self):
							for model_class_name in self.all_model_classes:
													__UpperCAmelCase			:							List[Any]	    =	model_class_name.from_pretrained('''albert-base-v2''')
													__UpperCAmelCase			:							Dict	    =	model(np.ones((1, 1)))
													self.assertIsNotNone(lowercase__)






@require_flax
class 							lowerCamelCase   (  unittest.TestCase ):


	@slow
	def  A(    self):
							__UpperCAmelCase			:							Union[str, Any]	    =	FlaxAlbertModel.from_pretrained('''albert-base-v2''')
							__UpperCAmelCase			:							Dict	    =	np.array([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]])
							__UpperCAmelCase			:							Tuple	    =	np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
							__UpperCAmelCase			:							Optional[int]	    =	model(lowercase__				,     attention_mask=lowercase__)[0]
							__UpperCAmelCase			:							int	    =	(1, 1_1, 7_6_8)
							self.assertEqual(output.shape				,     lowercase__)
							__UpperCAmelCase			:							Dict	    =	np.array(
							    [[[-0.6_5_1_3, 1.5_0_3_5, -0.2_7_6_6], [-0.6_5_1_5, 1.5_0_4_6, -0.2_7_8_0], [-0.6_5_1_2, 1.5_0_4_9, -0.2_7_8_4]]])

							self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]				,     lowercase__				,     atol=1e-4))





from typing import List, Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase       		=    logging.get_logger(__name__)

lowerCAmelCase       		=    {
    """huggingface/autoformer-tourism-monthly""": """https://huggingface.co/huggingface/autoformer-tourism-monthly/resolve/main/config.json""",
}




class 							lowerCamelCase   (  _UpperCamelCase ):
	_lowerCAmelCase     :		Optional[int]										=						'''autoformer'''
	_lowerCAmelCase     :		Any										=						{
	    '''hidden_size''': '''d_model''',
	    '''num_attention_heads''': '''encoder_attention_heads''',
	    '''num_hidden_layers''': '''encoder_layers''',
	}


	def __init__(    self				,     lowercase__ = None				,     lowercase__ = None				,     lowercase__ = "student_t"				,     lowercase__ = "nll"				,     lowercase__ = 1				,     lowercase__ = [1, 2, 3, 4, 5, 6, 7]				,     lowercase__ = True				,     lowercase__ = 0				,     lowercase__ = 0				,     lowercase__ = 0				,     lowercase__ = 0				,     lowercase__ = None				,     lowercase__ = None				,     lowercase__ = 6_4				,     lowercase__ = 2				,     lowercase__ = 2				,     lowercase__ = 2				,     lowercase__ = 2				,     lowercase__ = 3_2				,     lowercase__ = 3_2				,     lowercase__ = "gelu"				,     lowercase__ = 0.1				,     lowercase__ = 0.1				,     lowercase__ = 0.1				,     lowercase__ = 0.1				,     lowercase__ = 0.1				,     lowercase__ = 1_0_0				,     lowercase__ = 0.0_2				,     lowercase__ = True				,     lowercase__=True				,     lowercase__ = 1_0				,     lowercase__ = 2_5				,     lowercase__ = 3				,     **lowercase__				,     ):
							# time series specific configuration
							__UpperCAmelCase			:							Optional[int]	    =	prediction_length
							__UpperCAmelCase			:							Tuple	    =	context_length if context_length is not None else prediction_length
							__UpperCAmelCase			:							Tuple	    =	distribution_output
							__UpperCAmelCase			:							Tuple	    =	loss
							__UpperCAmelCase			:							Dict	    =	input_size
							__UpperCAmelCase			:							str	    =	num_time_features
							__UpperCAmelCase			:							str	    =	lags_sequence
							__UpperCAmelCase			:							int	    =	scaling
							__UpperCAmelCase			:							int	    =	num_dynamic_real_features
							__UpperCAmelCase			:							List[str]	    =	num_static_real_features
							__UpperCAmelCase			:							List[str]	    =	num_static_categorical_features
							if cardinality is not None and num_static_categorical_features > 0:
													if len(lowercase__) != num_static_categorical_features:
																			raise ValueError(
																			    '''The cardinality should be a list of the same length as `num_static_categorical_features`''')
													__UpperCAmelCase			:							Union[str, Any]	    =	cardinality
							else:
													__UpperCAmelCase			:							List[Any]	    =	[0]
							if embedding_dimension is not None and num_static_categorical_features > 0:
													if len(lowercase__) != num_static_categorical_features:
																			raise ValueError(
																			    '''The embedding dimension should be a list of the same length as `num_static_categorical_features`''')
													__UpperCAmelCase			:							Union[str, Any]	    =	embedding_dimension
							else:
													__UpperCAmelCase			:							int	    =	[min(5_0				,     (cat + 1) // 2) for cat in self.cardinality]
							__UpperCAmelCase			:							Tuple	    =	num_parallel_samples

							# Transformer architecture configuration
							__UpperCAmelCase			:							Dict	    =	input_size * len(self.lags_sequence) + self._number_of_features
							__UpperCAmelCase			:							Optional[int]	    =	d_model
							__UpperCAmelCase			:							List[Any]	    =	encoder_attention_heads
							__UpperCAmelCase			:							int	    =	decoder_attention_heads
							__UpperCAmelCase			:							str	    =	encoder_ffn_dim
							__UpperCAmelCase			:							int	    =	decoder_ffn_dim
							__UpperCAmelCase			:							List[Any]	    =	encoder_layers
							__UpperCAmelCase			:							Tuple	    =	decoder_layers

							__UpperCAmelCase			:							Union[str, Any]	    =	dropout
							__UpperCAmelCase			:							int	    =	attention_dropout
							__UpperCAmelCase			:							int	    =	activation_dropout
							__UpperCAmelCase			:							str	    =	encoder_layerdrop
							__UpperCAmelCase			:							Dict	    =	decoder_layerdrop

							__UpperCAmelCase			:							int	    =	activation_function
							__UpperCAmelCase			:							Optional[int]	    =	init_std

							__UpperCAmelCase			:							Optional[int]	    =	use_cache

							# Autoformer
							__UpperCAmelCase			:							str	    =	label_length
							__UpperCAmelCase			:							int	    =	moving_average
							__UpperCAmelCase			:							Optional[int]	    =	autocorrelation_factor

							super().__init__(is_encoder_decoder=lowercase__				,     **lowercase__)


	@property
	def  A(    self):
							return (
							    sum(self.embedding_dimension)
							    + self.num_dynamic_real_features
							    + self.num_time_features
							    + self.num_static_real_features
							    + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
							)

