





import tempfile
import unittest

import numpy as np
from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_flax


if is_flax_available():
     import os

     from flax.core.frozen_dict import unfreeze
     from flax.traverse_util import flatten_dict

     from transformers import FlaxBertModel

     __A							   =							"0.12"  # assumed parallelism: 8



@require_flax
@is_staging_test
class   _SCREAMING_SNAKE_CASE		(      unittest.TestCase       ):







 '''simple docstring'''

 @classmethod
 def 				SCREAMING_SNAKE_CASE_   (cls							:  str)			->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__: Optional[int]         =TOKEN
   HfFolder.save_token(UpperCAmelCase_)

 @classmethod
 def 				SCREAMING_SNAKE_CASE_   (cls							:  Optional[Any])			->int:



   '''simple docstring'''



   try:
     delete_repo(token=cls._token			,					repo_id="test-model-flax")
   except HTTPError:
     pass

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

 def 				SCREAMING_SNAKE_CASE_   (self							:  Any)			->Tuple:



   '''simple docstring'''



   lowerCamelCase__: str         =BertConfig(
       vocab_size=99			,					hidden_size=32			,					num_hidden_layers=5			,					num_attention_heads=4			,					intermediate_size=37)
   lowerCamelCase__: Dict         =FlaxBertModel(UpperCAmelCase_)
   model.push_to_hub("test-model-flax"			,					use_auth_token=self._token)

   lowerCamelCase__: Union[str, Any]         =FlaxBertModel.from_pretrained(F"""{USER}/test-model-flax""")

   lowerCamelCase__: str         =flatten_dict(unfreeze(model.params))
   lowerCamelCase__: Union[str, Any]         =flatten_dict(unfreeze(new_model.params))

   for key in base_params.keys():
     lowerCamelCase__: List[str]         =(base_params[key] - new_params[key]).sum().item()
     self.assertLessEqual(UpperCAmelCase_			,					1E-3			,					msg=F"""{key} not identical""")

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

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

   lowerCamelCase__: Tuple         =FlaxBertModel.from_pretrained(F"""{USER}/test-model-flax""")

   lowerCamelCase__: Optional[int]         =flatten_dict(unfreeze(model.params))
   lowerCamelCase__: List[str]         =flatten_dict(unfreeze(new_model.params))

   for key in base_params.keys():
     lowerCamelCase__: List[Any]         =(base_params[key] - new_params[key]).sum().item()
     self.assertLessEqual(UpperCAmelCase_			,					1E-3			,					msg=F"""{key} not identical""")







 def 				SCREAMING_SNAKE_CASE_   (self							:  Union[str, Any])			->str:



   '''simple docstring'''



   lowerCamelCase__: Any         =BertConfig(
       vocab_size=99			,					hidden_size=32			,					num_hidden_layers=5			,					num_attention_heads=4			,					intermediate_size=37)
   lowerCamelCase__: List[str]         =FlaxBertModel(UpperCAmelCase_)
   model.push_to_hub("valid_org/test-model-flax-org"			,					use_auth_token=self._token)

   lowerCamelCase__: Optional[int]         =FlaxBertModel.from_pretrained("valid_org/test-model-flax-org")

   lowerCamelCase__: int         =flatten_dict(unfreeze(model.params))
   lowerCamelCase__: Any         =flatten_dict(unfreeze(new_model.params))

   for key in base_params.keys():
     lowerCamelCase__: List[str]         =(base_params[key] - new_params[key]).sum().item()
     self.assertLessEqual(UpperCAmelCase_			,					1E-3			,					msg=F"""{key} not identical""")

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

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

   lowerCamelCase__: int         =FlaxBertModel.from_pretrained("valid_org/test-model-flax-org")

   lowerCamelCase__: Union[str, Any]         =flatten_dict(unfreeze(model.params))
   lowerCamelCase__: Union[str, Any]         =flatten_dict(unfreeze(new_model.params))

   for key in base_params.keys():
     lowerCamelCase__: Tuple         =(base_params[key] - new_params[key]).sum().item()
     self.assertLessEqual(UpperCAmelCase_			,					1E-3			,					msg=F"""{key} not identical""")





def 	lowerCAmelCase_					(			__a      ,       __a						)	->						Union[str, Any]:







  """simple docstring"""





  lowerCamelCase__: int         =True
  lowerCamelCase__: Any         =flatten_dict(modela.params						)
  lowerCamelCase__: Any         =flatten_dict(modela.params						)
  for key in flat_params_a.keys():
    if np.sum(np.abs(flat_params_a[key] - flat_params_a[key]						)						) > 1e-4:
      lowerCamelCase__: Dict         =False

  return models_are_equal






@require_flax
class   _SCREAMING_SNAKE_CASE		(      unittest.TestCase       ):







 '''simple docstring'''

 def 				SCREAMING_SNAKE_CASE_   (self							:  str)			->Any:



   '''simple docstring'''



   lowerCamelCase__: Optional[Any]         =BertConfig.from_pretrained("hf-internal-testing/tiny-bert-flax-only")
   lowerCamelCase__: str         =FlaxBertModel(UpperCAmelCase_)

   lowerCamelCase__: int         ="bert"
   with tempfile.TemporaryDirectory() as tmp_dir:
     model.save_pretrained(os.path.join(UpperCAmelCase_			,					UpperCAmelCase_))

     with self.assertRaises(UpperCAmelCase_):
       lowerCamelCase__: Tuple         =FlaxBertModel.from_pretrained(UpperCAmelCase_)

     lowerCamelCase__: Optional[int]         =FlaxBertModel.from_pretrained(UpperCAmelCase_			,					subfolder=UpperCAmelCase_)

   self.assertTrue(check_models_equal(UpperCAmelCase_			,					UpperCAmelCase_))

 def 				SCREAMING_SNAKE_CASE_   (self							:  Any)			->Any:



   '''simple docstring'''



   lowerCamelCase__: Union[str, Any]         =BertConfig.from_pretrained("hf-internal-testing/tiny-bert-flax-only")
   lowerCamelCase__: List[Any]         =FlaxBertModel(UpperCAmelCase_)

   lowerCamelCase__: int         ="bert"
   with tempfile.TemporaryDirectory() as tmp_dir:
     model.save_pretrained(os.path.join(UpperCAmelCase_			,					UpperCAmelCase_)			,					max_shard_size="10KB")

     with self.assertRaises(UpperCAmelCase_):
       lowerCamelCase__: Dict         =FlaxBertModel.from_pretrained(UpperCAmelCase_)

     lowerCamelCase__: str         =FlaxBertModel.from_pretrained(UpperCAmelCase_			,					subfolder=UpperCAmelCase_)

   self.assertTrue(check_models_equal(UpperCAmelCase_			,					UpperCAmelCase_))

 def 				SCREAMING_SNAKE_CASE_   (self							:  List[Any])			->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__: str         ="bert"
   lowerCamelCase__: Optional[Any]         ="hf-internal-testing/tiny-random-bert-subfolder"

   with self.assertRaises(UpperCAmelCase_):
     lowerCamelCase__: Union[str, Any]         =FlaxBertModel.from_pretrained(UpperCAmelCase_)

   lowerCamelCase__: Tuple         =FlaxBertModel.from_pretrained(UpperCAmelCase_			,					subfolder=UpperCAmelCase_)

   self.assertIsNotNone(UpperCAmelCase_)







 def 				SCREAMING_SNAKE_CASE_   (self							:  Dict)			->Union[str, Any]:



   '''simple docstring'''



   lowerCamelCase__: List[str]         ="bert"
   lowerCamelCase__: Optional[int]         ="hf-internal-testing/tiny-random-bert-sharded-subfolder"
   with self.assertRaises(UpperCAmelCase_):
     lowerCamelCase__: Dict         =FlaxBertModel.from_pretrained(UpperCAmelCase_)

   lowerCamelCase__: Any         =FlaxBertModel.from_pretrained(UpperCAmelCase_			,					subfolder=UpperCAmelCase_)

   self.assertIsNotNone(UpperCAmelCase_)




'''simple docstring'''







import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCamelCase									=     logging.get_logger(__name__)

__lowerCamelCase									=     {
    '''google/pix2struct-textcaps-base''': (
        '''https://huggingface.co/google/pix2struct-textcaps-base/resolve/main/config.json'''
    ),
}


class        A__   (   _snake_case      ):
    lowercase						       =		"pix2struct_text_model"
    lowercase						       =		["past_key_values"]
    lowercase						       =		{
        "hidden_size": "hidden_size",
        "num_attention_heads": "num_heads",
        "num_hidden_layers": "num_layers",
    }



    def __init__(       self		,					UpperCamelCase__=50244		,					UpperCamelCase__=768		,					UpperCamelCase__=64		,					UpperCamelCase__=2048		,					UpperCamelCase__=12		,					UpperCamelCase__=12		,					UpperCamelCase__=32		,					UpperCamelCase__=128		,					UpperCamelCase__=0.1		,					UpperCamelCase__=1e-6		,					UpperCamelCase__=1.0		,					UpperCamelCase__="gelu_new"		,					UpperCamelCase__=0		,					UpperCamelCase__=False		,					UpperCamelCase__=0		,					UpperCamelCase__=1		,					UpperCamelCase__=False		,					UpperCamelCase__=True		,					**UpperCamelCase__		,					)    ->   List[str]:





           '''simple docstring'''
           A_						       =   vocab_size
           A_						       =   hidden_size
           A_						       =   d_kv
           A_						       =   d_ff
           A_						       =   num_layers
           A_						       =   num_heads
           A_						       =   relative_attention_num_buckets
           A_						       =   relative_attention_max_distance
           A_						       =   dropout_rate
           A_						       =   layer_norm_epsilon
           A_						       =   initializer_factor
           A_						       =   use_cache

           A_						       =   eos_token_id
           A_						       =   decoder_start_token_id

           # for backwards compatibility
           A_						       =   dense_act_fn

           super().__init__(
               pad_token_id=UpperCamelCase__		,					eos_token_id=UpperCamelCase__		,					decoder_start_token_id=UpperCamelCase__		,					tie_word_embeddings=UpperCamelCase__		,					is_decoder=UpperCamelCase__		,					**UpperCamelCase__		,					)







    @classmethod
    def        snake_case_      (       cls		,					UpperCamelCase__		,					**UpperCamelCase__							)    ->   "PretrainedConfig":





           '''simple docstring'''
           cls._set_token_in_kwargs(UpperCamelCase__							)

           A_       ,      A_						       =   cls.get_config_dict(UpperCamelCase__		,					**UpperCamelCase__							)

           # get the text config dict if we are loading from Pix2StructConfig
           if config_dict.get("""model_type"""							) == "pix2struct":
                  A_						       =   config_dict["""text_config"""]

           if "model_type" in config_dict and hasattr(cls		,					"""model_type"""							) and config_dict["model_type"] != cls.model_type:
                  logger.warning(
                      f'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
                      f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''							)

           return cls.from_dict(UpperCamelCase__		,					**UpperCamelCase__							)


class        A__   (   _snake_case      ):
    lowercase						       =		"pix2struct_vision_model"



    def __init__(       self		,					UpperCamelCase__=768		,					UpperCamelCase__=768		,					UpperCamelCase__=2048		,					UpperCamelCase__=64		,					UpperCamelCase__=12		,					UpperCamelCase__=12		,					UpperCamelCase__="gelu_new"		,					UpperCamelCase__=1e-6		,					UpperCamelCase__=0.0		,					UpperCamelCase__=0.0		,					UpperCamelCase__=1e-1_0		,					UpperCamelCase__=1.0		,					UpperCamelCase__=4096		,					UpperCamelCase__=32		,					UpperCamelCase__=128		,					**UpperCamelCase__		,					)    ->   int:





           '''simple docstring'''
           super().__init__(**UpperCamelCase__							)

           A_						       =   hidden_size
           A_						       =   patch_embed_hidden_size
           A_						       =   d_ff
           A_						       =   dropout_rate
           A_						       =   num_hidden_layers
           A_						       =   num_attention_heads
           A_						       =   initializer_range
           A_						       =   initializer_factor
           A_						       =   attention_dropout
           A_						       =   layer_norm_eps
           A_						       =   dense_act_fn
           A_						       =   seq_len
           A_						       =   relative_attention_num_buckets
           A_						       =   relative_attention_max_distance
           A_						       =   d_kv







    @classmethod
    def        snake_case_      (       cls		,					UpperCamelCase__		,					**UpperCamelCase__							)    ->   "PretrainedConfig":





           '''simple docstring'''
           cls._set_token_in_kwargs(UpperCamelCase__							)

           A_       ,      A_						       =   cls.get_config_dict(UpperCamelCase__		,					**UpperCamelCase__							)

           # get the vision config dict if we are loading from Pix2StructConfig
           if config_dict.get("""model_type"""							) == "pix2struct":
                  A_						       =   config_dict["""vision_config"""]

           if "model_type" in config_dict and hasattr(cls		,					"""model_type"""							) and config_dict["model_type"] != cls.model_type:
                  logger.warning(
                      f'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
                      f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''							)

           return cls.from_dict(UpperCamelCase__		,					**UpperCamelCase__							)


class        A__   (   _snake_case      ):
    lowercase						       =		"pix2struct"
    lowercase						       =		True



    def __init__(       self		,					UpperCamelCase__=None		,					UpperCamelCase__=None		,					UpperCamelCase__=1.0		,					UpperCamelCase__=0.02		,					UpperCamelCase__=False		,					UpperCamelCase__=False		,					UpperCamelCase__=True		,					**UpperCamelCase__		,					)    ->   str:





           '''simple docstring'''
           super().__init__(tie_word_embeddings=UpperCamelCase__		,					is_encoder_decoder=UpperCamelCase__		,					**UpperCamelCase__							)

           if text_config is None:
                  A_						       =   {}
                  logger.info("""text_config is None. Initializing the Pix2StructTextConfig with default values."""							)

           if vision_config is None:
                  A_						       =   {}
                  logger.info("""vision_config is None. Initializing the Pix2StructVisionConfig with default values."""							)

           A_						       =   PixaStructTextConfig(**UpperCamelCase__							)
           A_						       =   PixaStructVisionConfig(**UpperCamelCase__							)

           A_						       =   self.text_config.decoder_start_token_id
           A_						       =   self.text_config.pad_token_id
           A_						       =   self.text_config.eos_token_id

           A_						       =   initializer_factor
           A_						       =   initializer_range

           A_						       =   self.initializer_range
           A_						       =   self.initializer_range

           A_						       =   is_vqa



    @classmethod
    def        snake_case_      (       cls		,					UpperCamelCase__		,					UpperCamelCase__		,					**UpperCamelCase__							)    ->   Optional[int]:





           '''simple docstring'''
           return cls(text_config=text_config.to_dict()		,					vision_config=vision_config.to_dict()		,					**UpperCamelCase__							)







    def        snake_case_      (       self							)    ->   List[Any]:





           '''simple docstring'''
           A_						       =   copy.deepcopy(self.__dict__							)
           A_						       =   self.text_config.to_dict()
           A_						       =   self.vision_config.to_dict()
           A_						       =   self.__class__.model_type
           return output
