



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 							__snake_case		(	unittest.TestCase		):




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




							def   UpperCAmelCase__     (			self	:			Tuple):
									lowerCAmelCase_     :	int							   =					ids_tensor([self.batch_size, self.seq_length]							,							self.vocab_size)

									lowerCAmelCase_     :	Union[str, Any]							   =					None
									if self.use_attention_mask:
											lowerCAmelCase_     :	Optional[int]							   =					random_attention_mask([self.batch_size, self.seq_length])

									lowerCAmelCase_     :	List[Any]							   =					None
									if self.use_token_type_ids:
											lowerCAmelCase_     :	str							   =					ids_tensor([self.batch_size, self.seq_length]							,							self.type_vocab_size)

									lowerCAmelCase_     :	Optional[Any]							   =					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   UpperCAmelCase__     (			self	:			Dict):
									lowerCAmelCase_     :	int							   =					self.prepare_config_and_inputs()
									lowerCAmelCase_   ,    lowerCAmelCase_   ,    lowerCAmelCase_   ,    lowerCAmelCase_     :	List[Any]							   =					config_and_inputs
									lowerCAmelCase_     :	List[str]							   =					{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
									return config, inputs_dict





@require_flax
class 							__snake_case		(	_SCREAMING_SNAKE_CASE		,unittest.TestCase		):
							_a      =					(
							    (
							        FlaxAlbertModel,
							        FlaxAlbertForPreTraining,
							        FlaxAlbertForMaskedLM,
							        FlaxAlbertForMultipleChoice,
							        FlaxAlbertForQuestionAnswering,
							        FlaxAlbertForSequenceClassification,
							        FlaxAlbertForTokenClassification,
							        FlaxAlbertForQuestionAnswering,
							    )
							    if is_flax_available()
							    else ()
							)




							def   UpperCAmelCase__     (			self	:			Dict):
									lowerCAmelCase_     :	int							   =					FlaxAlbertModelTester(self)







							@slow
							def   UpperCAmelCase__     (			self	:			Optional[int]):



									for model_class_name in self.all_model_classes:
											lowerCAmelCase_     :	Dict							   =					model_class_name.from_pretrained('''albert-base-v2''')
											lowerCAmelCase_     :	Tuple							   =					model(np.ones((1, 1)))
											self.assertIsNotNone(_lowercase)





@require_flax
class 							__snake_case		(	unittest.TestCase		):




							@slow
							def   UpperCAmelCase__     (			self	:			Dict):
									lowerCAmelCase_     :	Any							   =					FlaxAlbertModel.from_pretrained('''albert-base-v2''')
									lowerCAmelCase_     :	List[str]							   =					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]])
									lowerCAmelCase_     :	Tuple							   =					np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
									lowerCAmelCase_     :	str							   =					model(_lowercase							,							attention_mask=_lowercase)[0]
									lowerCAmelCase_     :	Any							   =					(1, 1_1, 7_6_8)
									self.assertEqual(output.shape							,							_lowercase)
									lowerCAmelCase_     :	List[str]							   =					np.array(
									    [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]])

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







'''simple docstring'''



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


_lowercase          =						logging.get_logger(__name__)

_lowercase          =						{
    """facebook/dpr-ctx_encoder-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-question_encoder-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-reader-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-ctx_encoder-multiset-base""": (
        """https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/config.json"""
    ),
    """facebook/dpr-question_encoder-multiset-base""": (
        """https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/config.json"""
    ),
    """facebook/dpr-reader-multiset-base""": (
        """https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/config.json"""
    ),
}




class     UpperCAmelCase_	(						_SCREAMING_SNAKE_CASE					):






						'''simple docstring'''






						_lowercase						: str										=  '''dpr'''




						def __init__(   self       ,  _lowercase=30_522       ,  _lowercase=768       ,  _lowercase=12       ,  _lowercase=12       ,  _lowercase=3_072       ,  _lowercase="gelu"       ,  _lowercase=0.1       ,  _lowercase=0.1       ,  _lowercase=512       ,  _lowercase=2       ,  _lowercase=0.02       ,  _lowercase=1e-12       ,  _lowercase=0       ,  _lowercase="absolute"       ,  _lowercase = 0       ,  **_lowercase       ,  ):

												"""simple docstring"""



												super().__init__(pad_token_id=_lowercase       ,  **_lowercase )

												_lowerCAmelCase    		=					vocab_size
												_lowerCAmelCase    		=					hidden_size
												_lowerCAmelCase    		=					num_hidden_layers
												_lowerCAmelCase    		=					num_attention_heads
												_lowerCAmelCase    		=					hidden_act
												_lowerCAmelCase    		=					intermediate_size
												_lowerCAmelCase    		=					hidden_dropout_prob
												_lowerCAmelCase    		=					attention_probs_dropout_prob
												_lowerCAmelCase    		=					max_position_embeddings
												_lowerCAmelCase    		=					type_vocab_size
												_lowerCAmelCase    		=					initializer_range
												_lowerCAmelCase    		=					layer_norm_eps
												_lowerCAmelCase    		=					projection_dim
												_lowerCAmelCase    		=					position_embedding_type
