



import unittest

import numpy as np

from transformers import DistilBertConfig, 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.distilbert.modeling_flax_distilbert import (
			    FlaxDistilBertForMaskedLM,
			    FlaxDistilBertForMultipleChoice,
			    FlaxDistilBertForQuestionAnswering,
			    FlaxDistilBertForSequenceClassification,
			    FlaxDistilBertForTokenClassification,
			    FlaxDistilBertModel,
			)



class 	lowerCamelCase      (						unittest.TestCase  ):







					'''simple docstring'''







					def __init__(			self    ,	lowerCAmelCase    ,	lowerCAmelCase=13    ,	lowerCAmelCase=7    ,	lowerCAmelCase=True    ,	lowerCAmelCase=True    ,	lowerCAmelCase=True    ,	lowerCAmelCase=True    ,	lowerCAmelCase=99    ,	lowerCAmelCase=32    ,	lowerCAmelCase=5    ,	lowerCAmelCase=4    ,	lowerCAmelCase=37    ,	lowerCAmelCase="gelu"    ,	lowerCAmelCase=0.1    ,	lowerCAmelCase=0.1    ,	lowerCAmelCase=512    ,	lowerCAmelCase=16    ,	lowerCAmelCase=2    ,	lowerCAmelCase=0.02    ,	lowerCAmelCase=4    ,	):
						UpperCAmelCase_						=				parent
						UpperCAmelCase_						=				batch_size
						UpperCAmelCase_						=				seq_length
						UpperCAmelCase_						=				is_training
						UpperCAmelCase_						=				use_attention_mask
						UpperCAmelCase_						=				use_token_type_ids
						UpperCAmelCase_						=				use_labels
						UpperCAmelCase_						=				vocab_size
						UpperCAmelCase_						=				hidden_size
						UpperCAmelCase_						=				num_hidden_layers
						UpperCAmelCase_						=				num_attention_heads
						UpperCAmelCase_						=				intermediate_size
						UpperCAmelCase_						=				hidden_act
						UpperCAmelCase_						=				hidden_dropout_prob
						UpperCAmelCase_						=				attention_probs_dropout_prob
						UpperCAmelCase_						=				max_position_embeddings
						UpperCAmelCase_						=				type_vocab_size
						UpperCAmelCase_						=				type_sequence_label_size
						UpperCAmelCase_						=				initializer_range
						UpperCAmelCase_						=				num_choices



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

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

						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    ,	tie_weights_=lowerCAmelCase    ,	)

						return config, input_ids, attention_mask



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





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







					'''simple docstring'''







					lowerCAmelCase_				:    List[Any]									=  (
					    (
					        FlaxDistilBertModel,
					        FlaxDistilBertForMaskedLM,
					        FlaxDistilBertForMultipleChoice,
					        FlaxDistilBertForQuestionAnswering,
					        FlaxDistilBertForSequenceClassification,
					        FlaxDistilBertForTokenClassification,
					        FlaxDistilBertForQuestionAnswering,
					    )
					    if is_flax_available()
					    else ()
					)



					def    A__		(			self     ):
						UpperCAmelCase_						=				FlaxDistilBertModelTester(self     )



					@slow
					def    A__		(			self     ):
						for model_class_name in self.all_model_classes:
							UpperCAmelCase_						=				model_class_name.from_pretrained("distilbert-base-uncased"     )
							UpperCAmelCase_						=				model(np.ones((1, 1)     )     )
							self.assertIsNotNone(lowerCAmelCase     )





@require_flax
class 	lowerCamelCase      (						unittest.TestCase  ):







					'''simple docstring'''







					@slow
					def    A__		(			self     ):
						UpperCAmelCase_						=				FlaxDistilBertModel.from_pretrained("distilbert-base-uncased"     )
						UpperCAmelCase_						=				np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]     )
						UpperCAmelCase_						=				np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]     )
						UpperCAmelCase_						=				model(lowerCAmelCase    ,	attention_mask=lowerCAmelCase     )[0]
						UpperCAmelCase_						=				(1, 11, 768)
						self.assertEqual(output.shape    ,	lowerCAmelCase     )
						UpperCAmelCase_						=				np.array([[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]     )

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






from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


SCREAMING_SNAKE_CASE     				=						logging.get_logger(__name__)

SCREAMING_SNAKE_CASE     				=						{
    "xlm-roberta-base": "https://huggingface.co/xlm-roberta-base/resolve/main/config.json",
    "xlm-roberta-large": "https://huggingface.co/xlm-roberta-large/resolve/main/config.json",
    "xlm-roberta-large-finetuned-conll02-dutch": (
        "https://huggingface.co/xlm-roberta-large-finetuned-conll02-dutch/resolve/main/config.json"
    ),
    "xlm-roberta-large-finetuned-conll02-spanish": (
        "https://huggingface.co/xlm-roberta-large-finetuned-conll02-spanish/resolve/main/config.json"
    ),
    "xlm-roberta-large-finetuned-conll03-english": (
        "https://huggingface.co/xlm-roberta-large-finetuned-conll03-english/resolve/main/config.json"
    ),
    "xlm-roberta-large-finetuned-conll03-german": (
        "https://huggingface.co/xlm-roberta-large-finetuned-conll03-german/resolve/main/config.json"
    ),
}



class 	lowerCamelCase      (						lowercase__  ):







					'''simple docstring'''







					lowerCAmelCase_				:    List[str]									=  'xlm-roberta'



					def __init__(			self    ,	lowerCAmelCase=3_0522    ,	lowerCAmelCase=768    ,	lowerCAmelCase=12    ,	lowerCAmelCase=12    ,	lowerCAmelCase=3072    ,	lowerCAmelCase="gelu"    ,	lowerCAmelCase=0.1    ,	lowerCAmelCase=0.1    ,	lowerCAmelCase=512    ,	lowerCAmelCase=2    ,	lowerCAmelCase=0.02    ,	lowerCAmelCase=1e-1_2    ,	lowerCAmelCase=1    ,	lowerCAmelCase=0    ,	lowerCAmelCase=2    ,	lowerCAmelCase="absolute"    ,	lowerCAmelCase=True    ,	lowerCAmelCase=None    ,	**lowerCAmelCase    ,	):
							super().__init__(pad_token_id=lowerCAmelCase    ,	bos_token_id=lowerCAmelCase    ,	eos_token_id=lowerCAmelCase    ,	**lowerCAmelCase     )

							UpperCAmelCase_						=				vocab_size
							UpperCAmelCase_						=				hidden_size
							UpperCAmelCase_						=				num_hidden_layers
							UpperCAmelCase_						=				num_attention_heads
							UpperCAmelCase_						=				hidden_act
							UpperCAmelCase_						=				intermediate_size
							UpperCAmelCase_						=				hidden_dropout_prob
							UpperCAmelCase_						=				attention_probs_dropout_prob
							UpperCAmelCase_						=				max_position_embeddings
							UpperCAmelCase_						=				type_vocab_size
							UpperCAmelCase_						=				initializer_range
							UpperCAmelCase_						=				layer_norm_eps
							UpperCAmelCase_						=				position_embedding_type
							UpperCAmelCase_						=				use_cache
							UpperCAmelCase_						=				classifier_dropout





class 	lowerCamelCase      (						lowercase__  ):







					'''simple docstring'''







					@property
					def    A__		(			self     ):
							if self.task == "multiple-choice":
									UpperCAmelCase_						=				{0: "batch", 1: "choice", 2: "sequence"}
							else:
									UpperCAmelCase_						=				{0: "batch", 1: "sequence"}
							return OrderedDict(
							    [
							        ("input_ids", dynamic_axis),
							        ("attention_mask", dynamic_axis),
							    ]     )

