






import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD




class       snake_case__ (				nn.Module):

       '''simple docstring'''
       lowerCamelCase				:					int
       lowerCamelCase				:					int
       lowerCamelCase				:					float								=			0.0
       lowerCamelCase				:					int								=			1
       lowerCamelCase				:					int								=			1
       lowerCamelCase				:					bool								=			True
       lowerCamelCase				:					bool								=			False
       lowerCamelCase				:					bool								=			False
       lowerCamelCase				:					bool								=			False
       lowerCamelCase				:					jnp.dtype								=			jnp.floataa



       def        __lowercase    (				self     )			->   int:






            '''simple docstring'''





            __snake_case      :Dict 							=    []
            __snake_case      :Tuple 							=    []

            for i in range(self.num_layers     ):
                 __snake_case      :Tuple 							=    self.in_channels if i == 0 else self.out_channels

                 __snake_case      :Optional[int] 							=    FlaxResnetBlockaD(
                     in_channels=a__      ,    out_channels=self.out_channels      ,    dropout_prob=self.dropout      ,    dtype=self.dtype      ,    )
                 resnets.append(a__     )

                 __snake_case      :Tuple 							=    FlaxTransformeraDModel(
                     in_channels=self.out_channels      ,    n_heads=self.num_attention_heads      ,    d_head=self.out_channels // self.num_attention_heads      ,    depth=1      ,    use_linear_projection=self.use_linear_projection      ,    only_cross_attention=self.only_cross_attention      ,    use_memory_efficient_attention=self.use_memory_efficient_attention      ,    dtype=self.dtype      ,    )
                 attentions.append(a__     )

            __snake_case      :Dict 							=    resnets
            __snake_case      :Any 							=    attentions

            if self.add_downsample:
                 __snake_case      :Optional[int] 							=    FlaxDownsampleaD(self.out_channels      ,    dtype=self.dtype     )



       def __call__(				self      ,    a__      ,    a__      ,    a__      ,    a__=True     )			->   int:






            '''simple docstring'''





            __snake_case      :Dict 							=    ()

            for resnet, attn in zip(self.resnets      ,    self.attentions     ):
                 __snake_case      :Optional[Any] 							=    resnet(a__      ,    a__      ,    deterministic=a__     )
                 __snake_case      :Any 							=    attn(a__      ,    a__      ,    deterministic=a__     )
                 output_states += (hidden_states,)

            if self.add_downsample:
                 __snake_case      :Optional[int] 							=    self.downsamplers_a(a__     )
                 output_states += (hidden_states,)

            return hidden_states, output_states




class       snake_case__ (				nn.Module):

       '''simple docstring'''
       lowerCamelCase				:					int
       lowerCamelCase				:					int
       lowerCamelCase				:					float								=			0.0
       lowerCamelCase				:					int								=			1
       lowerCamelCase				:					bool								=			True
       lowerCamelCase				:					jnp.dtype								=			jnp.floataa



       def        __lowercase    (				self     )			->   Optional[int]:






            '''simple docstring'''





            __snake_case      :Tuple 							=    []

            for i in range(self.num_layers     ):
                 __snake_case      :str 							=    self.in_channels if i == 0 else self.out_channels

                 __snake_case      :Dict 							=    FlaxResnetBlockaD(
                     in_channels=a__      ,    out_channels=self.out_channels      ,    dropout_prob=self.dropout      ,    dtype=self.dtype      ,    )
                 resnets.append(a__     )
            __snake_case      :List[Any] 							=    resnets

            if self.add_downsample:
                 __snake_case      :int 							=    FlaxDownsampleaD(self.out_channels      ,    dtype=self.dtype     )



       def __call__(				self      ,    a__      ,    a__      ,    a__=True     )			->   Tuple:






            '''simple docstring'''





            __snake_case      :Any 							=    ()

            for resnet in self.resnets:
                 __snake_case      :List[str] 							=    resnet(a__      ,    a__      ,    deterministic=a__     )
                 output_states += (hidden_states,)

            if self.add_downsample:
                 __snake_case      :Optional[int] 							=    self.downsamplers_a(a__     )
                 output_states += (hidden_states,)

            return hidden_states, output_states




class       snake_case__ (				nn.Module):

       '''simple docstring'''
       lowerCamelCase				:					int
       lowerCamelCase				:					int
       lowerCamelCase				:					int
       lowerCamelCase				:					float								=			0.0
       lowerCamelCase				:					int								=			1
       lowerCamelCase				:					int								=			1
       lowerCamelCase				:					bool								=			True
       lowerCamelCase				:					bool								=			False
       lowerCamelCase				:					bool								=			False
       lowerCamelCase				:					bool								=			False
       lowerCamelCase				:					jnp.dtype								=			jnp.floataa



       def        __lowercase    (				self     )			->   int:






            '''simple docstring'''





            __snake_case      :Dict 							=    []
            __snake_case      :Any 							=    []

            for i in range(self.num_layers     ):
                 __snake_case      :Tuple 							=    self.in_channels if (i == self.num_layers - 1) else self.out_channels
                 __snake_case      :Tuple 							=    self.prev_output_channel if i == 0 else self.out_channels

                 __snake_case      :Optional[int] 							=    FlaxResnetBlockaD(
                     in_channels=resnet_in_channels + res_skip_channels      ,    out_channels=self.out_channels      ,    dropout_prob=self.dropout      ,    dtype=self.dtype      ,    )
                 resnets.append(a__     )

                 __snake_case      :Optional[int] 							=    FlaxTransformeraDModel(
                     in_channels=self.out_channels      ,    n_heads=self.num_attention_heads      ,    d_head=self.out_channels // self.num_attention_heads      ,    depth=1      ,    use_linear_projection=self.use_linear_projection      ,    only_cross_attention=self.only_cross_attention      ,    use_memory_efficient_attention=self.use_memory_efficient_attention      ,    dtype=self.dtype      ,    )
                 attentions.append(a__     )

            __snake_case      :Union[str, Any] 							=    resnets
            __snake_case      :Tuple 							=    attentions

            if self.add_upsample:
                 __snake_case      :Any 							=    FlaxUpsampleaD(self.out_channels      ,    dtype=self.dtype     )



       def __call__(				self      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__=True     )			->   List[str]:






            '''simple docstring'''





            for resnet, attn in zip(self.resnets      ,    self.attentions     ):
                 # pop res hidden states
                 __snake_case      :int 							=    res_hidden_states_tuple[-1]
                 __snake_case      :Optional[int] 							=    res_hidden_states_tuple[:-1]
                 __snake_case      :int 							=    jnp.concatenate((hidden_states, res_hidden_states)      ,    axis=-1     )

                 __snake_case      :Dict 							=    resnet(a__      ,    a__      ,    deterministic=a__     )
                 __snake_case      :List[Any] 							=    attn(a__      ,    a__      ,    deterministic=a__     )

            if self.add_upsample:
                 __snake_case      :Optional[Any] 							=    self.upsamplers_a(a__     )

            return hidden_states




class       snake_case__ (				nn.Module):

       '''simple docstring'''
       lowerCamelCase				:					int
       lowerCamelCase				:					int
       lowerCamelCase				:					int
       lowerCamelCase				:					float								=			0.0
       lowerCamelCase				:					int								=			1
       lowerCamelCase				:					bool								=			True
       lowerCamelCase				:					jnp.dtype								=			jnp.floataa



       def        __lowercase    (				self     )			->   int:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    []

            for i in range(self.num_layers     ):
                 __snake_case      :int 							=    self.in_channels if (i == self.num_layers - 1) else self.out_channels
                 __snake_case      :List[str] 							=    self.prev_output_channel if i == 0 else self.out_channels

                 __snake_case      :Union[str, Any] 							=    FlaxResnetBlockaD(
                     in_channels=resnet_in_channels + res_skip_channels      ,    out_channels=self.out_channels      ,    dropout_prob=self.dropout      ,    dtype=self.dtype      ,    )
                 resnets.append(a__     )

            __snake_case      :Union[str, Any] 							=    resnets

            if self.add_upsample:
                 __snake_case      :Union[str, Any] 							=    FlaxUpsampleaD(self.out_channels      ,    dtype=self.dtype     )



       def __call__(				self      ,    a__      ,    a__      ,    a__      ,    a__=True     )			->   List[str]:






            '''simple docstring'''





            for resnet in self.resnets:
                 # pop res hidden states
                 __snake_case      :List[str] 							=    res_hidden_states_tuple[-1]
                 __snake_case      :Any 							=    res_hidden_states_tuple[:-1]
                 __snake_case      :Union[str, Any] 							=    jnp.concatenate((hidden_states, res_hidden_states)      ,    axis=-1     )

                 __snake_case      :Optional[Any] 							=    resnet(a__      ,    a__      ,    deterministic=a__     )

            if self.add_upsample:
                 __snake_case      :Tuple 							=    self.upsamplers_a(a__     )

            return hidden_states




class       snake_case__ (				nn.Module):

       '''simple docstring'''
       lowerCamelCase				:					int
       lowerCamelCase				:					float								=			0.0
       lowerCamelCase				:					int								=			1
       lowerCamelCase				:					int								=			1
       lowerCamelCase				:					bool								=			False
       lowerCamelCase				:					bool								=			False
       lowerCamelCase				:					jnp.dtype								=			jnp.floataa



       def        __lowercase    (				self     )			->   List[str]:






            '''simple docstring'''





            __snake_case      :Union[str, Any] 							=    [
                FlaxResnetBlockaD(
                    in_channels=self.in_channels      ,    out_channels=self.in_channels      ,    dropout_prob=self.dropout      ,    dtype=self.dtype      ,    )
            ]

            __snake_case      :str 							=    []

            for _ in range(self.num_layers     ):
                 __snake_case      :Dict 							=    FlaxTransformeraDModel(
                     in_channels=self.in_channels      ,    n_heads=self.num_attention_heads      ,    d_head=self.in_channels // self.num_attention_heads      ,    depth=1      ,    use_linear_projection=self.use_linear_projection      ,    use_memory_efficient_attention=self.use_memory_efficient_attention      ,    dtype=self.dtype      ,    )
                 attentions.append(a__     )

                 __snake_case      :int 							=    FlaxResnetBlockaD(
                     in_channels=self.in_channels      ,    out_channels=self.in_channels      ,    dropout_prob=self.dropout      ,    dtype=self.dtype      ,    )
                 resnets.append(a__     )

            __snake_case      :Optional[int] 							=    resnets
            __snake_case      :Union[str, Any] 							=    attentions



       def __call__(				self      ,    a__      ,    a__      ,    a__      ,    a__=True     )			->   int:






            '''simple docstring'''





            __snake_case      :Dict 							=    self.resnets[0](a__      ,    a__     )
            for attn, resnet in zip(self.attentions      ,    self.resnets[1:]     ):
                 __snake_case      :Tuple 							=    attn(a__      ,    a__      ,    deterministic=a__     )
                 __snake_case      :Dict 							=    resnet(a__      ,    a__      ,    deterministic=a__     )

            return hidden_states







"""simple docstring"""







import shutil
import tempfile
import unittest

from transformers import ClapFeatureExtractor, ClapProcessor, RobertaTokenizer, RobertaTokenizerFast
from transformers.testing_utils import require_sentencepiece, require_torchaudio

from .test_feature_extraction_clap import floats_list







@require_torchaudio
@require_sentencepiece
class 					A_(unittest.TestCase      ):







			"""simple docstring"""



			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					Any	       =    'laion/clap-htsat-unfused'
						_lowerCamelCase					:					Optional[Any]	       =    tempfile.mkdtemp()
			def    _lowerCAmelCase		(	self				,   **A							):
						return RobertaTokenizer.from_pretrained(self.checkpoint				,   **A							)
			def    _lowerCAmelCase		(	self				,   **A							):
						return ClapFeatureExtractor.from_pretrained(self.checkpoint				,   **A							)
			def    _lowerCAmelCase		(	self							):
						shutil.rmtree(self.tmpdirname							)
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					Optional[Any]	       =    self.get_tokenizer()
						_lowerCamelCase					:					Union[str, Any]	       =    self.get_feature_extractor()

						_lowerCamelCase					:					List[str]	       =    ClapProcessor(tokenizer=A				,   feature_extractor=A							)

						processor.save_pretrained(self.tmpdirname							)
						_lowerCamelCase					:					str	       =    ClapProcessor.from_pretrained(self.tmpdirname							)

						self.assertEqual(processor.tokenizer.get_vocab()				,   tokenizer.get_vocab()							)
						self.assertIsInstance(processor.tokenizer				,   A							)

						self.assertEqual(processor.feature_extractor.to_json_string()				,   feature_extractor.to_json_string()							)
						self.assertIsInstance(processor.feature_extractor				,   A							)
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					int	       =    ClapProcessor(tokenizer=self.get_tokenizer()				,   feature_extractor=self.get_feature_extractor()							)
						processor.save_pretrained(self.tmpdirname							)

						_lowerCamelCase					:					Tuple	       =    self.get_tokenizer(bos_token='(BOS)'				,   eos_token='(EOS)'							)
						_lowerCamelCase					:					str	       =    self.get_feature_extractor(do_normalize=A				,   padding_value=1.0							)

						_lowerCamelCase					:					List[Any]	       =    ClapProcessor.from_pretrained(
						    self.tmpdirname				,   bos_token='(BOS)'				,   eos_token='(EOS)'				,   do_normalize=A				,   padding_value=1.0							)

						self.assertEqual(processor.tokenizer.get_vocab()				,   tokenizer_add_kwargs.get_vocab()							)
						self.assertIsInstance(processor.tokenizer				,   A							)

						self.assertEqual(processor.feature_extractor.to_json_string()				,   feature_extractor_add_kwargs.to_json_string()							)
						self.assertIsInstance(processor.feature_extractor				,   A							)
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					Optional[int]	       =    self.get_feature_extractor()
						_lowerCamelCase					:					Optional[int]	       =    self.get_tokenizer()

						_lowerCamelCase					:					Any	       =    ClapProcessor(tokenizer=A				,   feature_extractor=A							)

						_lowerCamelCase					:					int	       =    floats_list((3, 1000)							)

						_lowerCamelCase					:					str	       =    feature_extractor(A				,   return_tensors='np'							)
						_lowerCamelCase					:					List[str]	       =    processor(audios=A				,   return_tensors='np'							)

						for key in input_feat_extract.keys():
									self.assertAlmostEqual(input_feat_extract[key].sum()				,   input_processor[key].sum()				,   delta=1E-2							)
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					Any	       =    self.get_feature_extractor()
						_lowerCamelCase					:					str	       =    self.get_tokenizer()

						_lowerCamelCase					:					List[Any]	       =    ClapProcessor(tokenizer=A				,   feature_extractor=A							)

						_lowerCamelCase					:					List[str]	       =    'This is a test string'

						_lowerCamelCase					:					Any	       =    processor(text=A							)

						_lowerCamelCase					:					Optional[int]	       =    tokenizer(A							)

						for key in encoded_tok.keys():
									self.assertListEqual(encoded_tok[key]				,   encoded_processor[key]							)
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					int	       =    self.get_feature_extractor()
						_lowerCamelCase					:					Union[str, Any]	       =    self.get_tokenizer()

						_lowerCamelCase					:					Optional[int]	       =    ClapProcessor(tokenizer=A				,   feature_extractor=A							)

						_lowerCamelCase					:					Dict	       =    [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

						_lowerCamelCase					:					int	       =    processor.batch_decode(A							)
						_lowerCamelCase					:					int	       =    tokenizer.batch_decode(A							)

						self.assertListEqual(A				,   A							)

			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase					:					Union[str, Any]	       =    self.get_feature_extractor()
						_lowerCamelCase					:					Optional[Any]	       =    self.get_tokenizer()

						_lowerCamelCase					:					int	       =    ClapProcessor(tokenizer=A				,   feature_extractor=A							)

						self.assertListEqual(
						    processor.model_input_names[2:]				,   feature_extractor.model_input_names				,   msg='`processor` and `feature_extractor` model input names do not match'				,   )





