





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







       """simple docstring"""





       @parameterized.expand([(None,), ('foo.json',)]	)
       def  __UpperCAmelCase			(     self						:						Any    ,__A						:						Any	)    ->       Optional[Any]:
         _lowercase						    =			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	)
           _lowercase						    =			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						:						int	)    ->       List[str]:
         _lowercase						    =			AutoConfig.from_pretrained('gpt2'	)
         _lowercase						    =			GenerationConfig.from_model_config(__A	)
         _lowercase						    =			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						:						Optional[int]	)    ->       Any:
         _lowercase						    =			GenerationConfig()
         _lowercase						    =			{
             'max_new_tokens': 1024,
             'foo': 'bar',
         }
         _lowercase						    =			copy.deepcopy(__A	)
         _lowercase						    =			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    ,1024	)

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







       def  __UpperCAmelCase			(     self						:						List[Any]	)    ->       List[Any]:
         _lowercase						    =			GenerationConfig()
         _lowercase						    =			'bar'

         with tempfile.TemporaryDirectory('test-generation-config'	) as tmp_dir:
           generation_config.save_pretrained(__A	)

           _lowercase						    =			GenerationConfig.from_pretrained(__A	)
         # update_kwargs was used to update the config on valid attributes
         self.assertEqual(new_config.foo    ,'bar'	)

         _lowercase						    =			GenerationConfig.from_model_config(__A	)
         assert not hasattr(__A    ,'foo'	)  # no new kwargs should be initialized if from config







       def  __UpperCAmelCase			(     self						:						List[Any]	)    ->       Dict:
         _lowercase						    =			GenerationConfig()
         self.assertEqual(default_config.temperature    ,1.0	)
         self.assertEqual(default_config.do_sample    ,__A	)
         self.assertEqual(default_config.num_beams    ,1	)

         _lowercase						    =			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	)
           _lowercase						    =			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 					A_					(   unittest.TestCase						):







       """simple docstring"""





       @classmethod
       def  __UpperCAmelCase			(     cls						:						Optional[Any]	)    ->       Optional[int]:
         _lowercase						    =			TOKEN
         HfFolder.save_token(__A	)







       @classmethod
       def  __UpperCAmelCase			(     cls						:						List[str]	)    ->       Union[str, 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						:						Tuple	)    ->       List[Any]:
         _lowercase						    =			GenerationConfig(
             do_sample=__A    ,temperature=0.7    ,length_penalty=1.0    ,)
         config.push_to_hub('test-generation-config'    ,use_auth_token=self._token	)

         _lowercase						    =			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	)

         _lowercase						    =			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						:						Optional[int]	)    ->       Tuple:
         _lowercase						    =			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	)

         _lowercase						    =			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	)

         _lowercase						    =			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 re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


__magic_name__   :  Optional[int]        =							logging.get_logger(__name__)
def 		lowercase__  (				_UpperCamelCase) ->			Dict:
      """simple docstring"""

      UpperCamelCase					       =  r'\w+[.]\d+'
      UpperCamelCase					       =  re.findall(_UpperCamelCase       ,							_UpperCamelCase)
      for pat in pats:
            UpperCamelCase					       =  key.replace(_UpperCamelCase       ,							'_'.join(pat.split('.')))
      return key
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase) ->			Tuple:
      """simple docstring"""

      UpperCamelCase					       =  pt_tuple_key[:-1] + ('scale',)
      if (
          any('norm' in str_ for str_ in pt_tuple_key)
          and (pt_tuple_key[-1] == "bias")
          and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
          and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
      ):
            UpperCamelCase					       =  pt_tuple_key[:-1] + ('scale',)
            return renamed_pt_tuple_key, pt_tensor
      elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
            UpperCamelCase					       =  pt_tuple_key[:-1] + ('scale',)
            return renamed_pt_tuple_key, pt_tensor

      # embedding
      if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
            UpperCamelCase					       =  pt_tuple_key[:-1] + ('embedding',)
            return renamed_pt_tuple_key, pt_tensor

      # conv layer
      UpperCamelCase					       =  pt_tuple_key[:-1] + ('kernel',)
      if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
            UpperCamelCase					       =  pt_tensor.transpose(2       ,							3       ,							1       ,							0)
            return renamed_pt_tuple_key, pt_tensor

      # linear layer
      UpperCamelCase					       =  pt_tuple_key[:-1] + ('kernel',)
      if pt_tuple_key[-1] == "weight":
            UpperCamelCase					       =  pt_tensor.T
            return renamed_pt_tuple_key, pt_tensor

      # old PyTorch layer norm weight
      UpperCamelCase					       =  pt_tuple_key[:-1] + ('weight',)
      if pt_tuple_key[-1] == "gamma":
            return renamed_pt_tuple_key, pt_tensor

      # old PyTorch layer norm bias
      UpperCamelCase					       =  pt_tuple_key[:-1] + ('bias',)
      if pt_tuple_key[-1] == "beta":
            return renamed_pt_tuple_key, pt_tensor

      return pt_tuple_key, pt_tensor
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase=42) ->			Optional[Any]:
      """simple docstring"""

      UpperCamelCase					       =  {k: v.numpy() for k, v in pt_state_dict.items()}

      # Step 2: Since the model is stateless, get random Flax params
      UpperCamelCase					       =  flax_model.init_weights(PRNGKey(_UpperCamelCase))

      UpperCamelCase					       =  flatten_dict(_UpperCamelCase)
      UpperCamelCase					       =  {}

      # Need to change some parameters name to match Flax names
      for pt_key, pt_tensor in pt_state_dict.items():
            UpperCamelCase					       =  rename_key(_UpperCamelCase)
            UpperCamelCase					       =  tuple(renamed_pt_key.split('.'))

            # Correctly rename weight parameters
            UpperCamelCase       ,  UpperCamelCase					       =  rename_key_and_reshape_tensor(_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase)

            if flax_key in random_flax_state_dict:
                  if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
                        raise ValueError(
                            F'PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '
                            F'{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.')

        # also add unexpected weight so that warning is thrown
            UpperCamelCase					       =  jnp.asarray(_UpperCamelCase)

      return unflatten_dict(_UpperCamelCase)



