'''simple docstring'''




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 							a__(   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			,       ):




					"""simple docstring"""


					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 				a_			(  self):




					"""simple docstring"""


					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 				a_			(  self):




					"""simple docstring"""


					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 							a__(   lowerCAmelCase__					,		unittest.TestCase      ):


			'''simple docstring'''






			UpperCAmelCase_     :	int						=			(
			    (
			        FlaxAlbertModel,
			        FlaxAlbertForPreTraining,
			        FlaxAlbertForMaskedLM,
			        FlaxAlbertForMultipleChoice,
			        FlaxAlbertForQuestionAnswering,
			        FlaxAlbertForSequenceClassification,
			        FlaxAlbertForTokenClassification,
			        FlaxAlbertForQuestionAnswering,
			    )
			    if is_flax_available()
			    else ()
			)



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    FlaxAlbertModelTester(self)







			@slow
			def 				a_			(  self):




					"""simple docstring"""


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


			'''simple docstring'''






			@slow
			def 				a_			(  self):




					"""simple docstring"""


					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.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]			,       __lowerCAmelCase			,       atol=1E-4))





'''simple docstring'''




import copy
from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Audio, Features, Value
from .base import TaskTemplate






@dataclass(frozen=lowerCAmelCase__      )
class 							a__(   lowerCAmelCase__      ):


			'''simple docstring'''






			UpperCAmelCase_     :	str						=			field(default='''automatic-speech-recognition'''					,		metadata={'''include_in_asdict_even_if_is_default''': True}      )
			UpperCAmelCase_     :	ClassVar[Features]						=			Features({'''audio''': Audio()}      )
			UpperCAmelCase_     :	ClassVar[Features]						=			Features({'''transcription''': Value('''string'''      )}      )
			UpperCAmelCase_     :	str						=			"audio"
			UpperCAmelCase_     :	str						=			"transcription"



			def 				a_			(  self			,       __lowerCAmelCase):




					"""simple docstring"""


					if self.audio_column not in features:
							raise ValueError(f"Column {self.audio_column} is not present in features.")
					if not isinstance(features[self.audio_column]			,       __lowerCAmelCase):
							raise ValueError(f"Column {self.audio_column} is not an Audio type.")
					lowerCAmelCase     =    copy.deepcopy(self)
					lowerCAmelCase     =    self.input_schema.copy()
					lowerCAmelCase     =    features[self.audio_column]
					lowerCAmelCase     =    input_schema
					return task_template







			@property
			def 				a_			(  self):




					"""simple docstring"""


					return {self.audio_column: "audio", self.transcription_column: "transcription"}




