
import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotConfig, is_flax_available
from transformers.testing_utils import jax_device, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
			import os

			# The slow tests are often failing with OOM error on GPU
			# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
			# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
			_A						      =				'platform'

			import jax
			import jax.numpy as jnp

			from transformers import BlenderbotTokenizer
			from transformers.models.blenderbot.modeling_flax_blenderbot import (
			    FlaxBlenderbotForConditionalGeneration,
			    FlaxBlenderbotModel,
			    shift_tokens_right,
			)



def 		lowercase_    (     A__				,	A__				,	A__=None				,	A__=None				,	A__=None				,	A__=None				,	A__=None				,	A__=None				,	)							->							int:
				"""simple docstring"""

				if attention_mask is None:
								snake_case				  =      np.where(input_ids != config.pad_token_id				,	1				,	0	)
				if decoder_attention_mask is None:
								snake_case				  =      np.where(decoder_input_ids != config.pad_token_id				,	1				,	0	)
				if head_mask is None:
								snake_case				  =      np.ones((config.encoder_layers, config.encoder_attention_heads)	)
				if decoder_head_mask is None:
								snake_case				  =      np.ones((config.decoder_layers, config.decoder_attention_heads)	)
				if cross_attn_head_mask is None:
								snake_case				  =      np.ones((config.decoder_layers, config.decoder_attention_heads)	)
				return {
				    "input_ids": input_ids,
				    "decoder_input_ids": decoder_input_ids,
				    "attention_mask": attention_mask,
				    "decoder_attention_mask": attention_mask,
				}







class   lowerCamelCase	:





					def __init__(self   :					List[str]       ,					_A   :					Any       ,					_A   :					Optional[Any]=1_3       ,					_A   :					Optional[Any]=7       ,					_A   :					Optional[Any]=True       ,					_A   :					Any=False       ,					_A   :					Union[str, Any]=9_9       ,					_A   :					Dict=1_6       ,					_A   :					Optional[Any]=2       ,					_A   :					str=4       ,					_A   :					Dict=4       ,					_A   :					List[str]="gelu"       ,					_A   :					int=0.1       ,					_A   :					Optional[int]=0.1       ,					_A   :					int=3_2       ,					_A   :					List[Any]=2       ,					_A   :					Any=1       ,					_A   :					List[Any]=0       ,					_A   :					List[Any]=0.02       ,					)		->		Tuple:
									snake_case				  =      parent
									snake_case				  =      batch_size
									snake_case				  =      seq_length
									snake_case				  =      is_training
									snake_case				  =      use_labels
									snake_case				  =      vocab_size
									snake_case				  =      hidden_size
									snake_case				  =      num_hidden_layers
									snake_case				  =      num_attention_heads
									snake_case				  =      intermediate_size
									snake_case				  =      hidden_act
									snake_case				  =      hidden_dropout_prob
									snake_case				  =      attention_probs_dropout_prob
									snake_case				  =      max_position_embeddings
									snake_case				  =      eos_token_id
									snake_case				  =      pad_token_id
									snake_case				  =      bos_token_id
									snake_case				  =      initializer_range





					def 	UpperCAmelCase(self   :					Any     )		->		Optional[Any]:
									snake_case				  =      np.clip(ids_tensor([self.batch_size, self.seq_length - 1]       ,					self.vocab_size     )       ,					3       ,					self.vocab_size     )
									snake_case				  =      np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)       ,					dtype=np.intaa     ))       ,					-1     )

									snake_case				  =      shift_tokens_right(_SCREAMING_SNAKE_CASE       ,					1       ,					2     )

									snake_case				  =      BlenderbotConfig(
									    vocab_size=self.vocab_size       ,					d_model=self.hidden_size       ,					encoder_layers=self.num_hidden_layers       ,					decoder_layers=self.num_hidden_layers       ,					encoder_attention_heads=self.num_attention_heads       ,					decoder_attention_heads=self.num_attention_heads       ,					encoder_ffn_dim=self.intermediate_size       ,					decoder_ffn_dim=self.intermediate_size       ,					dropout=self.hidden_dropout_prob       ,					attention_dropout=self.attention_probs_dropout_prob       ,					max_position_embeddings=self.max_position_embeddings       ,					eos_token_id=self.eos_token_id       ,					bos_token_id=self.bos_token_id       ,					pad_token_id=self.pad_token_id       ,					initializer_range=self.initializer_range       ,					use_cache=_SCREAMING_SNAKE_CASE       ,					)
									snake_case				  =      prepare_blenderbot_inputs_dict(_SCREAMING_SNAKE_CASE       ,					_SCREAMING_SNAKE_CASE       ,					_SCREAMING_SNAKE_CASE     )
									return config, inputs_dict





					def 	UpperCAmelCase(self   :					int     )		->		Dict:
									snake_case  ,							snake_case				  =      self.prepare_config_and_inputs()
									return config, inputs_dict





					def 	UpperCAmelCase(self   :					str       ,					_A   :					Optional[Any]       ,					_A   :					Any       ,					_A   :					List[str]     )		->		Optional[int]:
									snake_case				  =      2_0
									snake_case				  =      model_class_name(_SCREAMING_SNAKE_CASE     )

									snake_case				  =      model.encode(inputs_dict["input_ids"]     )

									snake_case  ,							snake_case				  =      (
									    inputs_dict["decoder_input_ids"],
									    inputs_dict["decoder_attention_mask"],
									)

									snake_case				  =      model.init_cache(decoder_input_ids.shape[0]       ,					_SCREAMING_SNAKE_CASE       ,					_SCREAMING_SNAKE_CASE     )
									snake_case				  =      jnp.ones((decoder_input_ids.shape[0], max_decoder_length)       ,					dtype="i4"     )

									snake_case				  =      jnp.broadcast_to(
									    jnp.arange(decoder_input_ids.shape[-1] - 1     )[None, :]       ,					(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)       ,					)
									snake_case				  =      model.decode(
									    decoder_input_ids[:, :-1]       ,					_SCREAMING_SNAKE_CASE       ,					decoder_attention_mask=_SCREAMING_SNAKE_CASE       ,					past_key_values=_SCREAMING_SNAKE_CASE       ,					decoder_position_ids=_SCREAMING_SNAKE_CASE       ,					)

									snake_case				  =      jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]       ,					dtype="i4"     )
									snake_case				  =      model.decode(
									    decoder_input_ids[:, -1:]       ,					_SCREAMING_SNAKE_CASE       ,					decoder_attention_mask=_SCREAMING_SNAKE_CASE       ,					past_key_values=outputs_cache.past_key_values       ,					decoder_position_ids=_SCREAMING_SNAKE_CASE       ,					)

									snake_case				  =      model.decode(_SCREAMING_SNAKE_CASE       ,					_SCREAMING_SNAKE_CASE     )

									snake_case				  =      np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])     )     )
									self.parent.assertTrue(diff < 1E-3       ,					msg=f'Max diff is {diff}'     )





					def 	UpperCAmelCase(self   :					Dict       ,					_A   :					Optional[int]       ,					_A   :					Tuple       ,					_A   :					List[Any]     )		->		List[Any]:
									snake_case				  =      2_0
									snake_case				  =      model_class_name(_SCREAMING_SNAKE_CASE     )

									snake_case				  =      model.encode(inputs_dict["input_ids"]     )

									snake_case  ,							snake_case				  =      (
									    inputs_dict["decoder_input_ids"],
									    inputs_dict["decoder_attention_mask"],
									)

									snake_case				  =      jnp.concatenate(
									    [
									        decoder_attention_mask,
									        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])     ),
									    ]       ,					axis=-1       ,					)

									snake_case				  =      model.init_cache(decoder_input_ids.shape[0]       ,					_SCREAMING_SNAKE_CASE       ,					_SCREAMING_SNAKE_CASE     )
									snake_case				  =      jnp.broadcast_to(
									    jnp.arange(decoder_input_ids.shape[-1] - 1     )[None, :]       ,					(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)       ,					)

									snake_case				  =      model.decode(
									    decoder_input_ids[:, :-1]       ,					_SCREAMING_SNAKE_CASE       ,					decoder_attention_mask=_SCREAMING_SNAKE_CASE       ,					past_key_values=_SCREAMING_SNAKE_CASE       ,					decoder_position_ids=_SCREAMING_SNAKE_CASE       ,					)
									snake_case				  =      jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]       ,					dtype="i4"     )
									snake_case				  =      model.decode(
									    decoder_input_ids[:, -1:]       ,					_SCREAMING_SNAKE_CASE       ,					past_key_values=outputs_cache.past_key_values       ,					decoder_attention_mask=_SCREAMING_SNAKE_CASE       ,					decoder_position_ids=_SCREAMING_SNAKE_CASE       ,					)

									snake_case				  =      model.decode(_SCREAMING_SNAKE_CASE       ,					_SCREAMING_SNAKE_CASE       ,					decoder_attention_mask=_SCREAMING_SNAKE_CASE     )

									snake_case				  =      np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])     )     )
									self.parent.assertTrue(diff < 1E-3       ,					msg=f'Max diff is {diff}'     )







@require_flax
class   lowerCamelCase	(						unittest.TestCase						):
					UpperCAmelCase__       :							Union[str, Any]     =	99





					def 	UpperCAmelCase(self   :					List[str]     )		->		Dict:
									snake_case				  =      np.array(
									    [
									        [7_1, 8_2, 1_8, 3_3, 4_6, 9_1, 2],
									        [6_8, 3_4, 2_6, 5_8, 3_0, 8_2, 2],
									        [5, 9_7, 1_7, 3_9, 9_4, 4_0, 2],
									        [7_6, 8_3, 9_4, 2_5, 7_0, 7_8, 2],
									        [8_7, 5_9, 4_1, 3_5, 4_8, 6_6, 2],
									        [5_5, 1_3, 1_6, 5_8, 5, 2, 1],  # note padding
									        [6_4, 2_7, 3_1, 5_1, 1_2, 7_5, 2],
									        [5_2, 6_4, 8_6, 1_7, 8_3, 3_9, 2],
									        [4_8, 6_1, 9, 2_4, 7_1, 8_2, 2],
									        [2_6, 1, 6_0, 4_8, 2_2, 1_3, 2],
									        [2_1, 5, 6_2, 2_8, 1_4, 7_6, 2],
									        [4_5, 9_8, 3_7, 8_6, 5_9, 4_8, 2],
									        [7_0, 7_0, 5_0, 9, 2_8, 0, 2],
									    ]       ,					dtype=np.intaa       ,					)

									snake_case				  =      input_ids.shape[0]
									snake_case				  =      BlenderbotConfig(
									    vocab_size=self.vocab_size       ,					d_model=2_4       ,					encoder_layers=2       ,					decoder_layers=2       ,					encoder_attention_heads=2       ,					decoder_attention_heads=2       ,					encoder_ffn_dim=3_2       ,					decoder_ffn_dim=3_2       ,					max_position_embeddings=4_8       ,					eos_token_id=2       ,					pad_token_id=1       ,					bos_token_id=0       ,					)
									return config, input_ids, batch_size





					def 	UpperCAmelCase(self   :					str     )		->		Union[str, Any]:
									snake_case  ,							snake_case  ,							snake_case				  =      self._get_config_and_data()
									snake_case				  =      FlaxBlenderbotForConditionalGeneration(_SCREAMING_SNAKE_CASE     )
									snake_case				  =      lm_model(input_ids=_SCREAMING_SNAKE_CASE     )
									snake_case				  =      (batch_size, input_ids.shape[1], config.vocab_size)
									self.assertEqual(outputs["logits"].shape       ,					_SCREAMING_SNAKE_CASE     )





					def 	UpperCAmelCase(self   :					Optional[int]     )		->		Tuple:
									snake_case				  =      BlenderbotConfig(
									    vocab_size=self.vocab_size       ,					d_model=1_4       ,					encoder_layers=2       ,					decoder_layers=2       ,					encoder_attention_heads=2       ,					decoder_attention_heads=2       ,					encoder_ffn_dim=8       ,					decoder_ffn_dim=8       ,					max_position_embeddings=4_8       ,					)
									snake_case				  =      FlaxBlenderbotForConditionalGeneration(_SCREAMING_SNAKE_CASE     )
									snake_case				  =      np.array([[7_1, 8_2, 1_8, 3_3, 4_6, 9_1, 2], [6_8, 3_4, 2_6, 5_8, 3_0, 2, 1]]       ,					dtype=np.intaa     )
									snake_case				  =      np.array([[8_2, 7_1, 8_2, 1_8, 2], [5_8, 6_8, 2, 1, 1]]       ,					dtype=np.intaa     )
									snake_case				  =      lm_model(input_ids=_SCREAMING_SNAKE_CASE       ,					decoder_input_ids=_SCREAMING_SNAKE_CASE     )
									snake_case				  =      (*summary.shape, config.vocab_size)
									self.assertEqual(outputs["logits"].shape       ,					_SCREAMING_SNAKE_CASE     )





					def 	UpperCAmelCase(self   :					int     )		->		Optional[int]:
									snake_case				  =      np.array([[7_1, 8_2, 1_8, 3_3, 2, 1, 1], [6_8, 3_4, 2_6, 5_8, 3_0, 8_2, 2]]       ,					dtype=np.intaa     )
									snake_case				  =      shift_tokens_right(_SCREAMING_SNAKE_CASE       ,					1       ,					2     )
									snake_case				  =      np.equal(_SCREAMING_SNAKE_CASE       ,					1     ).astype(np.floataa     ).sum()
									snake_case				  =      np.equal(_SCREAMING_SNAKE_CASE       ,					1     ).astype(np.floataa     ).sum()
									self.assertEqual(shifted.shape       ,					input_ids.shape     )
									self.assertEqual(_SCREAMING_SNAKE_CASE       ,					n_pad_before - 1     )
									self.assertTrue(np.equal(shifted[:, 0]       ,					2     ).all()     )







@require_flax
class   lowerCamelCase	(						lowerCAmelCase__							,				unittest.TestCase							,				lowerCAmelCase__						):
					UpperCAmelCase__       :							List[Any]     =	True
					UpperCAmelCase__       :							List[Any]     =	(
					    (
					        FlaxBlenderbotModel,
					        FlaxBlenderbotForConditionalGeneration,
					    )
					    if is_flax_available()
					    else ()
					)
					UpperCAmelCase__       :							Dict     =	(FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()





					def 	UpperCAmelCase(self   :					Union[str, Any]     )		->		Any:
									snake_case				  =      FlaxBlenderbotModelTester(self     )





					def 	UpperCAmelCase(self   :					Any     )		->		Optional[int]:
									snake_case  ,							snake_case				  =      self.model_tester.prepare_config_and_inputs()
									for model_class in self.all_model_classes:
													self.model_tester.check_use_cache_forward(_SCREAMING_SNAKE_CASE       ,					_SCREAMING_SNAKE_CASE       ,					_SCREAMING_SNAKE_CASE     )





					def 	UpperCAmelCase(self   :					Union[str, Any]     )		->		Any:
									snake_case  ,							snake_case				  =      self.model_tester.prepare_config_and_inputs()
									for model_class in self.all_model_classes:
													self.model_tester.check_use_cache_forward_with_attn_mask(_SCREAMING_SNAKE_CASE       ,					_SCREAMING_SNAKE_CASE       ,					_SCREAMING_SNAKE_CASE     )





					def 	UpperCAmelCase(self   :					Optional[int]     )		->		int:
									snake_case  ,							snake_case				  =      self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
													with self.subTest(model_class.__name__     ):
																	snake_case				  =      self._prepare_for_class(_SCREAMING_SNAKE_CASE       ,					_SCREAMING_SNAKE_CASE     )
																	snake_case				  =      model_class(_SCREAMING_SNAKE_CASE     )

																	@jax.jit
																	def encode_jitted(_A   :					Any       ,					_A   :					List[str]=None       ,					**_A   :					str     ):
																					return model.encode(input_ids=_SCREAMING_SNAKE_CASE       ,					attention_mask=_SCREAMING_SNAKE_CASE     )

																	with self.subTest("JIT Enabled"     ):
																					snake_case				  =      encode_jitted(**_SCREAMING_SNAKE_CASE     ).to_tuple()

																	with self.subTest("JIT Disabled"     ):
																					with jax.disable_jit():
																									snake_case				  =      encode_jitted(**_SCREAMING_SNAKE_CASE     ).to_tuple()

																	self.assertEqual(len(_SCREAMING_SNAKE_CASE     )       ,					len(_SCREAMING_SNAKE_CASE     )     )
																	for jitted_output, output in zip(_SCREAMING_SNAKE_CASE       ,					_SCREAMING_SNAKE_CASE     ):
																					self.assertEqual(jitted_output.shape       ,					output.shape     )





					def 	UpperCAmelCase(self   :					Dict     )		->		str:
									snake_case  ,							snake_case				  =      self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
													with self.subTest(model_class.__name__     ):
																	snake_case				  =      model_class(_SCREAMING_SNAKE_CASE     )
																	snake_case				  =      model.encode(inputs_dict["input_ids"]       ,					inputs_dict["attention_mask"]     )

																	snake_case				  =      {
																	    "decoder_input_ids": inputs_dict["decoder_input_ids"],
																	    "decoder_attention_mask": inputs_dict["decoder_attention_mask"],
																	    "encoder_outputs": encoder_outputs,
																	}

																	@jax.jit
																	def decode_jitted(_A   :					Optional[Any]       ,					_A   :					int       ,					_A   :					int     ):
																					return model.decode(
																					    decoder_input_ids=_SCREAMING_SNAKE_CASE       ,					decoder_attention_mask=_SCREAMING_SNAKE_CASE       ,					encoder_outputs=_SCREAMING_SNAKE_CASE       ,					)

																	with self.subTest("JIT Enabled"     ):
																					snake_case				  =      decode_jitted(**_SCREAMING_SNAKE_CASE     ).to_tuple()

																	with self.subTest("JIT Disabled"     ):
																					with jax.disable_jit():
																									snake_case				  =      decode_jitted(**_SCREAMING_SNAKE_CASE     ).to_tuple()

																	self.assertEqual(len(_SCREAMING_SNAKE_CASE     )       ,					len(_SCREAMING_SNAKE_CASE     )     )
																	for jitted_output, output in zip(_SCREAMING_SNAKE_CASE       ,					_SCREAMING_SNAKE_CASE     ):
																					self.assertEqual(jitted_output.shape       ,					output.shape     )





					@slow
					def 	UpperCAmelCase(self   :					Optional[Any]     )		->		List[str]:
									for model_class_name in self.all_model_classes:
													snake_case				  =      model_class_name.from_pretrained("facebook/blenderbot-400M-distill"     )
													# FlaxBlenderbotForSequenceClassification expects eos token in input_ids
													snake_case				  =      np.ones((1, 1)     ) * model.config.eos_token_id
													snake_case				  =      model(_SCREAMING_SNAKE_CASE     )
													self.assertIsNotNone(_SCREAMING_SNAKE_CASE     )





					@unittest.skipUnless(jax_device != "cpu"       ,					"3B test too slow on CPU."     )
					@slow
					def 	UpperCAmelCase(self   :					str     )		->		Dict:
									snake_case				  =      {"num_beams": 1, "early_stopping": True, "min_length": 1_5, "max_length": 2_5}
									snake_case				  =      {"skip_special_tokens": True, "clean_up_tokenization_spaces": True}

									snake_case				  =      FlaxBlenderbotForConditionalGeneration.from_pretrained("facebook/blenderbot-3B"       ,					from_pt=_SCREAMING_SNAKE_CASE     )
									snake_case				  =      BlenderbotTokenizer.from_pretrained("facebook/blenderbot-3B"     )

									snake_case				  =      ["Sam"]
									snake_case				  =      tokenizer(_SCREAMING_SNAKE_CASE       ,					return_tensors="jax"     )

									snake_case				  =      model.generate(**_SCREAMING_SNAKE_CASE       ,					**_SCREAMING_SNAKE_CASE     )
									snake_case				  =      "Sam is a great name. It means \"sun\" in Gaelic."

									snake_case				  =      tokenizer.batch_decode(_SCREAMING_SNAKE_CASE       ,					**_SCREAMING_SNAKE_CASE     )
									assert generated_txt[0].strip() == tgt_text



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


_A						      =				{
    "configuration_wav2vec2": ["WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP", "Wav2Vec2Config"],
    "feature_extraction_wav2vec2": ["Wav2Vec2FeatureExtractor"],
    "processing_wav2vec2": ["Wav2Vec2Processor"],
    "tokenization_wav2vec2": ["Wav2Vec2CTCTokenizer", "Wav2Vec2Tokenizer"],
}


try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _A						      =				[
       "WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST",
       "Wav2Vec2ForAudioFrameClassification",
       "Wav2Vec2ForCTC",
       "Wav2Vec2ForMaskedLM",
       "Wav2Vec2ForPreTraining",
       "Wav2Vec2ForSequenceClassification",
       "Wav2Vec2ForXVector",
       "Wav2Vec2Model",
       "Wav2Vec2PreTrainedModel",
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _A						      =				[
       "TF_WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST",
       "TFWav2Vec2ForCTC",
       "TFWav2Vec2Model",
       "TFWav2Vec2PreTrainedModel",
       "TFWav2Vec2ForSequenceClassification",
   ]

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _A						      =				[
       "FlaxWav2Vec2ForCTC",
       "FlaxWav2Vec2ForPreTraining",
       "FlaxWav2Vec2Model",
       "FlaxWav2Vec2PreTrainedModel",
   ]


if TYPE_CHECKING:
   from .configuration_wavaveca import WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP, WavaVecaConfig
   from .feature_extraction_wavaveca import WavaVecaFeatureExtractor
   from .processing_wavaveca import WavaVecaProcessor
   from .tokenization_wavaveca import WavaVecaCTCTokenizer, WavaVecaTokenizer

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_wavaveca import (
          WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST,
          WavaVecaForAudioFrameClassification,
          WavaVecaForCTC,
          WavaVecaForMaskedLM,
          WavaVecaForPreTraining,
          WavaVecaForSequenceClassification,
          WavaVecaForXVector,
          WavaVecaModel,
          WavaVecaPreTrainedModel,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_wavaveca import (
          TF_WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFWavaVecaForCTC,
          TFWavaVecaForSequenceClassification,
          TFWavaVecaModel,
          TFWavaVecaPreTrainedModel,
      )

   try:
      if not is_flax_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_wavaveca import (
          FlaxWavaVecaForCTC,
          FlaxWavaVecaForPreTraining,
          FlaxWavaVecaModel,
          FlaxWavaVecaPreTrainedModel,
      )


else:
   import sys

   _A						      =				_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

