





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 SCREAMING_SNAKE_CASE_ (					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	,  ):
										lowerCamelCase__       =       parent
										lowerCamelCase__       =       batch_size
										lowerCamelCase__       =       seq_length
										lowerCamelCase__       =       is_training
										lowerCamelCase__       =       use_attention_mask
										lowerCamelCase__       =       use_token_type_ids
										lowerCamelCase__       =       use_labels
										lowerCamelCase__       =       vocab_size
										lowerCamelCase__       =       hidden_size
										lowerCamelCase__       =       num_hidden_layers
										lowerCamelCase__       =       num_attention_heads
										lowerCamelCase__       =       intermediate_size
										lowerCamelCase__       =       hidden_act
										lowerCamelCase__       =       hidden_dropout_prob
										lowerCamelCase__       =       attention_probs_dropout_prob
										lowerCamelCase__       =       max_position_embeddings
										lowerCamelCase__       =       type_vocab_size
										lowerCamelCase__       =       type_sequence_label_size
										lowerCamelCase__       =       initializer_range
										lowerCamelCase__       =       num_choices





					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       ids_tensor([self.batch_size, self.seq_length]	,  self.vocab_size	)

										lowerCamelCase__       =       None
										if self.use_attention_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__       =       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=_lowerCAmelCase	,  initializer_range=self.initializer_range	,  )

										return config, input_ids, token_type_ids, attention_mask






					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       self.prepare_config_and_inputs()
										lowerCamelCase__    ,		lowerCamelCase__    ,		lowerCamelCase__    ,		lowerCamelCase__       =       config_and_inputs
										lowerCamelCase__       =       {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
										return config, inputs_dict

@require_flax
class SCREAMING_SNAKE_CASE_ (					lowercase_							,       unittest.TestCase   ):





					"""simple docstring"""


					A__				 =    (
					    (
					        FlaxAlbertModel,
					        FlaxAlbertForPreTraining,
					        FlaxAlbertForMaskedLM,
					        FlaxAlbertForMultipleChoice,
					        FlaxAlbertForQuestionAnswering,
					        FlaxAlbertForSequenceClassification,
					        FlaxAlbertForTokenClassification,
					        FlaxAlbertForQuestionAnswering,
					    )
					    if is_flax_available()
					    else ()
					)





					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       FlaxAlbertModelTester(self	)






					@slow
					def 					__magic_name__ ( self	):
										for model_class_name in self.all_model_classes:
															lowerCamelCase__       =       model_class_name.from_pretrained("albert-base-v2"	)
															lowerCamelCase__       =       model(np.ones((1, 1)	)	)
															self.assertIsNotNone(_lowerCAmelCase	)

@require_flax
class SCREAMING_SNAKE_CASE_ (					unittest.TestCase   ):





					"""simple docstring"""


					@slow
					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       FlaxAlbertModel.from_pretrained("albert-base-v2"	)
										lowerCamelCase__       =       np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]	)
										lowerCamelCase__       =       np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]	)
										lowerCamelCase__       =       model(_lowerCAmelCase	,  attention_mask=_lowerCAmelCase	)[0]
										lowerCamelCase__       =       (1, 11, 768)
										self.assertEqual(output.shape	,  _lowerCAmelCase	)
										lowerCamelCase__       =       np.array(
										    [[[-0.6_513, 1.5_035, -0.2_766], [-0.6_515, 1.5_046, -0.2_780], [-0.6_512, 1.5_049, -0.2_784]]]	)

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






from collections import OrderedDict
from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Union

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


if TYPE_CHECKING:
						from ...onnx.config import PatchingSpec
						from ...tokenization_utils_base import PreTrainedTokenizerBase


A_												=     logging.get_logger(__name__)

A_												=     {
    "allenai/longformer-base-4096": "https://huggingface.co/allenai/longformer-base-4096/resolve/main/config.json",
    "allenai/longformer-large-4096": "https://huggingface.co/allenai/longformer-large-4096/resolve/main/config.json",
    "allenai/longformer-large-4096-finetuned-triviaqa": (
        "https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/config.json"
    ),
    "allenai/longformer-base-4096-extra.pos.embd.only": (
        "https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/config.json"
    ),
    "allenai/longformer-large-4096-extra.pos.embd.only": (
        "https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/config.json"
    ),
}

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					A__				 =    "longformer"





					def __init__( self	,  _lowerCAmelCase = 512	,  _lowerCAmelCase = 2	,  _lowerCAmelCase = 1	,  _lowerCAmelCase = 0	,  _lowerCAmelCase = 2	,  _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-12	,  _lowerCAmelCase = False	,  **_lowerCAmelCase	,  ):
										super().__init__(pad_token_id=_lowerCAmelCase	,  **_lowerCAmelCase	)

										lowerCamelCase__       =       attention_window
										lowerCamelCase__       =       sep_token_id
										lowerCamelCase__       =       bos_token_id
										lowerCamelCase__       =       eos_token_id
										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__       =       onnx_export

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase = "default"	,  _lowerCAmelCase = None	):
										super().__init__(_lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	)
										lowerCamelCase__       =       True





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





					@property
					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       super().outputs
										if self.task == "default":
															lowerCamelCase__       =       {0: "batch"}
										return outputs





					@property
					def 					__magic_name__ ( self	):
										return 1E-4





					@property
					def 					__magic_name__ ( self	):
										# needs to be >= 14 to support tril operator
										return max(super().default_onnx_opset	,  14	)






					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase = -1	,  _lowerCAmelCase = -1	,  _lowerCAmelCase = False	,  _lowerCAmelCase = None	,  ):
										lowerCamelCase__       =       super().generate_dummy_inputs(
										    preprocessor=_lowerCAmelCase	,  batch_size=_lowerCAmelCase	,  seq_length=_lowerCAmelCase	,  is_pair=_lowerCAmelCase	,  framework=_lowerCAmelCase	)
										import torch

										# for some reason, replacing this code by inputs["global_attention_mask"] = torch.randint(2, inputs["input_ids"].shape, dtype=torch.int64)
										# makes the export fail randomly
										lowerCamelCase__       =       torch.zeros_like(inputs["input_ids"]	)
										# make every second token global
										lowerCamelCase__       =       1

										return inputs

