





'''simple docstring'''







import copy
import tempfile
import unittest

from huggingface_hub import HfFolder, delete_repo
from parameterized import parameterized
from requests.exceptions import HTTPError

from transformers import AutoConfig, GenerationConfig
from transformers.testing_utils import TOKEN, USER, is_staging_test
class 				__SCREAMING_SNAKE_CASE  (unittest.TestCase      ):






				"""simple docstring"""







				@parameterized.expand([(None,), ("foo.json",)]					)
				def   UpperCamelCase__   (  self     :					Tuple					,				__a     :					Optional[Any]					):
								_a								= GenerationConfig(
								    do_sample=__a					,				temperature=0.7					,				length_penalty=1.0					,				bad_words_ids=[[1, 2, 3], [4, 5]]					,				)
								with tempfile.TemporaryDirectory() as tmp_dir:
												config.save_pretrained(__a					,				config_name=__a					)
												_a								= GenerationConfig.from_pretrained(__a					,				config_name=__a					)

								# Checks parameters that were specified
								self.assertEqual(loaded_config.do_sample					,				__a					)
								self.assertEqual(loaded_config.temperature					,				0.7					)
								self.assertEqual(loaded_config.length_penalty					,				1.0					)
								self.assertEqual(loaded_config.bad_words_ids					,				[[1, 2, 3], [4, 5]]					)

								# Checks parameters that were not specified (defaults)
								self.assertEqual(loaded_config.top_k					,				50					)
								self.assertEqual(loaded_config.max_length					,				20					)
								self.assertEqual(loaded_config.max_time					,				__a					)




				def   UpperCamelCase__   (  self     :					Union[str, Any]					):
								_a								= AutoConfig.from_pretrained("gpt2"					)
								_a								= GenerationConfig.from_model_config(__a					)
								_a								= GenerationConfig()

								# The generation config has loaded a few non-default parameters from the model config
								self.assertNotEqual(__a					,				__a					)

								# One of those parameters is eos_token_id -- check if it matches
								self.assertNotEqual(generation_config_from_model.eos_token_id					,				default_generation_config.eos_token_id					)
								self.assertEqual(generation_config_from_model.eos_token_id					,				model_config.eos_token_id					)




				def   UpperCamelCase__   (  self     :					int					):
								_a								= GenerationConfig()
								_a								= {
								    "max_new_tokens": 10_24,
								    "foo": "bar",
								}
								_a								= copy.deepcopy(__a					)
								_a								= generation_config.update(**__a					)

								# update_kwargs was not modified (no side effects)
								self.assertEqual(__a					,				__a					)

								# update_kwargs was used to update the config on valid attributes
								self.assertEqual(generation_config.max_new_tokens					,				10_24					)

								# `.update()` returns a dictionary of unused kwargs
								self.assertEqual(__a					,				{"foo": "bar"}					)




				def   UpperCamelCase__   (  self     :					Dict					):
								_a								= GenerationConfig()
								_a								= "bar"

								with tempfile.TemporaryDirectory("test-generation-config"					) as tmp_dir:
												generation_config.save_pretrained(__a					)

												_a								= GenerationConfig.from_pretrained(__a					)
								# update_kwargs was used to update the config on valid attributes
								self.assertEqual(new_config.foo					,				"bar"					)

								_a								= GenerationConfig.from_model_config(__a					)
								assert not hasattr(__a					,				"foo"					)  # no new kwargs should be initialized if from config




				def   UpperCamelCase__   (  self     :					Dict					):
								_a								= GenerationConfig()
								self.assertEqual(default_config.temperature					,				1.0					)
								self.assertEqual(default_config.do_sample					,				__a					)
								self.assertEqual(default_config.num_beams					,				1					)

								_a								= GenerationConfig(
								    do_sample=__a					,				temperature=0.7					,				length_penalty=1.0					,				bad_words_ids=[[1, 2, 3], [4, 5]]					,				)
								self.assertEqual(config.temperature					,				0.7					)
								self.assertEqual(config.do_sample					,				__a					)
								self.assertEqual(config.num_beams					,				1					)

								with tempfile.TemporaryDirectory() as tmp_dir:
												config.save_pretrained(__a					)
												_a								= GenerationConfig.from_pretrained(__a					,				temperature=1.0					)

								self.assertEqual(loaded_config.temperature					,				1.0					)
								self.assertEqual(loaded_config.do_sample					,				__a					)
								self.assertEqual(loaded_config.num_beams					,				1					)  # default value


@is_staging_test
class 				__SCREAMING_SNAKE_CASE  (unittest.TestCase      ):






				"""simple docstring"""







				@classmethod
				def   UpperCamelCase__   (  cls     :					Union[str, Any]					):
								_a								= TOKEN
								HfFolder.save_token(__a					)




				@classmethod
				def   UpperCamelCase__   (  cls     :					Optional[Any]					):
								try:
												delete_repo(token=cls._token					,				repo_id="test-generation-config"					)
								except HTTPError:
												pass

								try:
												delete_repo(token=cls._token					,				repo_id="valid_org/test-generation-config-org"					)
								except HTTPError:
												pass




				def   UpperCamelCase__   (  self     :					Optional[Any]					):
								_a								= GenerationConfig(
								    do_sample=__a					,				temperature=0.7					,				length_penalty=1.0					,				)
								config.push_to_hub("test-generation-config"					,				use_auth_token=self._token					)

								_a								= GenerationConfig.from_pretrained(f'{USER}/test-generation-config'					)
								for k, v in config.to_dict().items():
												if k != "transformers_version":
																self.assertEqual(__a					,				getattr(__a					,				__a					)					)

        # Reset repo
								delete_repo(token=self._token					,				repo_id="test-generation-config"					)

								# Push to hub via save_pretrained
								with tempfile.TemporaryDirectory() as tmp_dir:
												config.save_pretrained(
												    __a					,				repo_id="test-generation-config"					,				push_to_hub=__a					,				use_auth_token=self._token					)

								_a								= GenerationConfig.from_pretrained(f'{USER}/test-generation-config'					)
								for k, v in config.to_dict().items():
												if k != "transformers_version":
																self.assertEqual(__a					,				getattr(__a					,				__a					)					)




				def   UpperCamelCase__   (  self     :					int					):
								_a								= GenerationConfig(
								    do_sample=__a					,				temperature=0.7					,				length_penalty=1.0					,				)
								config.push_to_hub("valid_org/test-generation-config-org"					,				use_auth_token=self._token					)

								_a								= GenerationConfig.from_pretrained("valid_org/test-generation-config-org"					)
								for k, v in config.to_dict().items():
												if k != "transformers_version":
																self.assertEqual(__a					,				getattr(__a					,				__a					)					)

        # Reset repo
								delete_repo(token=self._token					,				repo_id="valid_org/test-generation-config-org"					)

								# Push to hub via save_pretrained
								with tempfile.TemporaryDirectory() as tmp_dir:
												config.save_pretrained(
												    __a					,				repo_id="valid_org/test-generation-config-org"					,				push_to_hub=__a					,				use_auth_token=self._token					)

								_a								= GenerationConfig.from_pretrained("valid_org/test-generation-config-org"					)
								for k, v in config.to_dict().items():
												if k != "transformers_version":
																self.assertEqual(__a					,				getattr(__a					,				__a					)					)




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

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




class A  (   nn.Module     ):
      UpperCamelCase__							:  int
      UpperCamelCase__							:  int
      UpperCamelCase__							:  float      =0.0
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  bool      =True
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  jnp.dtype      =jnp.floataa






      def        lowerCamelCase       (					self						:     Any )					->						Any:
             """simple docstring"""




             _lowerCamelCase       :		str		=[]
             _lowerCamelCase       :		Dict		=[]

             for i in range(self.num_layers ):
                    _lowerCamelCase       :		Union[str, Any]		=self.in_channels if i == 0 else self.out_channels

                    _lowerCamelCase       :		Any		=FlaxResnetBlockaD(
                        in_channels=lowercase_			,   out_channels=self.out_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
                    resnets.append(lowercase_ )

                    _lowerCamelCase       :		Dict		=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(lowercase_ )

             _lowerCamelCase       :		Optional[int]		=resnets
             _lowerCamelCase       :		Dict		=attentions

             if self.add_downsample:
                    _lowerCamelCase       :		Union[str, Any]		=FlaxDownsampleaD(self.out_channels			,   dtype=self.dtype )






      def __call__(					self						:     Any			,   lowercase_						:     Optional[Any]			,   lowercase_						:     List[Any]			,   lowercase_						:     Union[str, Any]			,   lowercase_						:     List[str]=True )					->						Union[str, Any]:
             """simple docstring"""




             _lowerCamelCase       :		Tuple		=()

             for resnet, attn in zip(self.resnets			,   self.attentions ):
                    _lowerCamelCase       :		Union[str, Any]		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    _lowerCamelCase       :		Union[str, Any]		=attn(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    output_states += (hidden_states,)

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

             return hidden_states, output_states




class A  (   nn.Module     ):
      UpperCamelCase__							:  int
      UpperCamelCase__							:  int
      UpperCamelCase__							:  float      =0.0
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  bool      =True
      UpperCamelCase__							:  jnp.dtype      =jnp.floataa






      def        lowerCamelCase       (					self						:     Optional[int] )					->						int:
             """simple docstring"""




             _lowerCamelCase       :		str		=[]

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

                    _lowerCamelCase       :		Any		=FlaxResnetBlockaD(
                        in_channels=lowercase_			,   out_channels=self.out_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
                    resnets.append(lowercase_ )
             _lowerCamelCase       :		Union[str, Any]		=resnets

             if self.add_downsample:
                    _lowerCamelCase       :		List[str]		=FlaxDownsampleaD(self.out_channels			,   dtype=self.dtype )






      def __call__(					self						:     Any			,   lowercase_						:     Dict			,   lowercase_						:     Union[str, Any]			,   lowercase_						:     Optional[Any]=True )					->						Dict:
             """simple docstring"""




             _lowerCamelCase       :		Optional[int]		=()

             for resnet in self.resnets:
                    _lowerCamelCase       :		Tuple		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    output_states += (hidden_states,)

             if self.add_downsample:
                    _lowerCamelCase       :		Tuple		=self.downsamplers_a(lowercase_ )
                    output_states += (hidden_states,)

             return hidden_states, output_states




class A  (   nn.Module     ):
      UpperCamelCase__							:  int
      UpperCamelCase__							:  int
      UpperCamelCase__							:  int
      UpperCamelCase__							:  float      =0.0
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  bool      =True
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  jnp.dtype      =jnp.floataa






      def        lowerCamelCase       (					self						:     Dict )					->						Tuple:
             """simple docstring"""




             _lowerCamelCase       :		str		=[]
             _lowerCamelCase       :		List[str]		=[]

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

                    _lowerCamelCase       :		List[str]		=FlaxResnetBlockaD(
                        in_channels=resnet_in_channels + res_skip_channels			,   out_channels=self.out_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
                    resnets.append(lowercase_ )

                    _lowerCamelCase       :		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(lowercase_ )

             _lowerCamelCase       :		int		=resnets
             _lowerCamelCase       :		Dict		=attentions

             if self.add_upsample:
                    _lowerCamelCase       :		str		=FlaxUpsampleaD(self.out_channels			,   dtype=self.dtype )






      def __call__(					self						:     Optional[Any]			,   lowercase_						:     Any			,   lowercase_						:     Any			,   lowercase_						:     List[Any]			,   lowercase_						:     Dict			,   lowercase_						:     Union[str, Any]=True )					->						Optional[int]:
             """simple docstring"""




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

                    _lowerCamelCase       :		Optional[Any]		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    _lowerCamelCase       :		List[Any]		=attn(lowercase_			,   lowercase_			,   deterministic=lowercase_ )

             if self.add_upsample:
                    _lowerCamelCase       :		Optional[Any]		=self.upsamplers_a(lowercase_ )

             return hidden_states




class A  (   nn.Module     ):
      UpperCamelCase__							:  int
      UpperCamelCase__							:  int
      UpperCamelCase__							:  int
      UpperCamelCase__							:  float      =0.0
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  bool      =True
      UpperCamelCase__							:  jnp.dtype      =jnp.floataa






      def        lowerCamelCase       (					self						:     List[Any] )					->						Dict:
             """simple docstring"""




             _lowerCamelCase       :		List[str]		=[]

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

                    _lowerCamelCase       :		str		=FlaxResnetBlockaD(
                        in_channels=resnet_in_channels + res_skip_channels			,   out_channels=self.out_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
                    resnets.append(lowercase_ )

             _lowerCamelCase       :		str		=resnets

             if self.add_upsample:
                    _lowerCamelCase       :		List[str]		=FlaxUpsampleaD(self.out_channels			,   dtype=self.dtype )






      def __call__(					self						:     List[str]			,   lowercase_						:     str			,   lowercase_						:     Dict			,   lowercase_						:     Union[str, Any]			,   lowercase_						:     Any=True )					->						int:
             """simple docstring"""




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

                    _lowerCamelCase       :		Optional[Any]		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )

             if self.add_upsample:
                    _lowerCamelCase       :		Union[str, Any]		=self.upsamplers_a(lowercase_ )

             return hidden_states




class A  (   nn.Module     ):
      UpperCamelCase__							:  int
      UpperCamelCase__							:  float      =0.0
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  int      =1
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  bool      =False
      UpperCamelCase__							:  jnp.dtype      =jnp.floataa






      def        lowerCamelCase       (					self						:     int )					->						Tuple:
             """simple docstring"""




             _lowerCamelCase       :		Optional[Any]		=[
                 FlaxResnetBlockaD(
                     in_channels=self.in_channels			,   out_channels=self.in_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
             ]

             _lowerCamelCase       :		Any		=[]

             for _ in range(self.num_layers ):
                    _lowerCamelCase       :		Optional[int]		=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(lowercase_ )

                    _lowerCamelCase       :		Tuple		=FlaxResnetBlockaD(
                        in_channels=self.in_channels			,   out_channels=self.in_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
                    resnets.append(lowercase_ )

             _lowerCamelCase       :		List[Any]		=resnets
             _lowerCamelCase       :		List[str]		=attentions






      def __call__(					self						:     Optional[int]			,   lowercase_						:     List[Any]			,   lowercase_						:     Dict			,   lowercase_						:     Tuple			,   lowercase_						:     List[str]=True )					->						int:
             """simple docstring"""




             _lowerCamelCase       :		Dict		=self.resnets[0](lowercase_			,   lowercase_ )
             for attn, resnet in zip(self.attentions			,   self.resnets[1:] ):
                    _lowerCamelCase       :		Tuple		=attn(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    _lowerCamelCase       :		List[str]		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )

             return hidden_states


