
"""simple docstring"""



import json
import os
import shutil
import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

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

from transformers import AutoConfig, BertConfig, GPTaConfig
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import TOKEN, USER, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / 'utils'))

from test_module.custom_configuration import CustomConfig  # noqa E402


UpperCAmelCase__          = {
    'return_dict': False,
    'output_hidden_states': True,
    'output_attentions': True,
    'torchscript': True,
    'torch_dtype': 'float16',
    'use_bfloat16': True,
    'tf_legacy_loss': True,
    'pruned_heads': {'a': 1},
    'tie_word_embeddings': False,
    'is_decoder': True,
    'cross_attention_hidden_size': 128,
    'add_cross_attention': True,
    'tie_encoder_decoder': True,
    'max_length': 50,
    'min_length': 3,
    'do_sample': True,
    'early_stopping': True,
    'num_beams': 3,
    'num_beam_groups': 3,
    'diversity_penalty': 0.5,
    'temperature': 2.0,
    'top_k': 10,
    'top_p': 0.7,
    'typical_p': 0.2,
    'repetition_penalty': 0.8,
    'length_penalty': 0.8,
    'no_repeat_ngram_size': 5,
    'encoder_no_repeat_ngram_size': 5,
    'bad_words_ids': [1, 2, 3],
    'num_return_sequences': 3,
    'chunk_size_feed_forward': 5,
    'output_scores': True,
    'return_dict_in_generate': True,
    'forced_bos_token_id': 2,
    'forced_eos_token_id': 3,
    'remove_invalid_values': True,
    'architectures': ['BertModel'],
    'finetuning_task': 'translation',
    'id2label': {0: 'label'},
    'label2id': {'label': '0'},
    'tokenizer_class': 'BertTokenizerFast',
    'prefix': 'prefix',
    'bos_token_id': 6,
    'pad_token_id': 7,
    'eos_token_id': 8,
    'sep_token_id': 9,
    'decoder_start_token_id': 10,
    'exponential_decay_length_penalty': (5, 1.01),
    'suppress_tokens': [0, 1],
    'begin_suppress_tokens': 2,
    'task_specific_params': {'translation': 'some_params'},
    'problem_type': 'regression',
}

@is_staging_test
class       lowerCAmelCase__				(						unittest.TestCase	):







     @classmethod
     def      lowercase	(							cls			:			Optional[Any]			):
       _snake_case					     =       TOKEN
       HfFolder.save_token(_lowerCamelCase			)







     @classmethod
     def      lowercase	(							cls			:			Optional[int]			):
       try:
         delete_repo(token=cls._token    ,      repo_id='''test-config'''			)
       except HTTPError:
         pass

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

       try:
         delete_repo(token=cls._token    ,      repo_id='''test-dynamic-config'''			)
       except HTTPError:
         pass







     def      lowercase	(							self			:			Any			):
       _snake_case					     =       BertConfig(
           vocab_size=99    ,      hidden_size=32    ,      num_hidden_layers=5    ,      num_attention_heads=4    ,      intermediate_size=37			)
       config.push_to_hub('''test-config'''    ,      use_auth_token=self._token			)

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

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

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

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







     def      lowercase	(							self			:			Optional[int]			):
       _snake_case					     =       BertConfig(
           vocab_size=99    ,      hidden_size=32    ,      num_hidden_layers=5    ,      num_attention_heads=4    ,      intermediate_size=37			)
       config.push_to_hub('''valid_org/test-config-org'''    ,      use_auth_token=self._token			)

       _snake_case					     =       BertConfig.from_pretrained('''valid_org/test-config-org'''			)
       for k, v in config.to_dict().items():
         if k != "transformers_version":
           self.assertEqual(_lowerCamelCase    ,      getattr(_lowerCamelCase    ,      _lowerCamelCase			)			)

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

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

       _snake_case					     =       BertConfig.from_pretrained('''valid_org/test-config-org'''			)
       for k, v in config.to_dict().items():
         if k != "transformers_version":
           self.assertEqual(_lowerCamelCase    ,      getattr(_lowerCamelCase    ,      _lowerCamelCase			)			)







     def      lowercase	(							self			:			Optional[Any]			):
       CustomConfig.register_for_auto_class()
       _snake_case					     =       CustomConfig(attribute=42			)

       config.push_to_hub('''test-dynamic-config'''    ,      use_auth_token=self._token			)

       # This has added the proper auto_map field to the config
       self.assertDictEqual(config.auto_map    ,      {'''AutoConfig''': '''custom_configuration.CustomConfig'''}			)

       _snake_case					     =       AutoConfig.from_pretrained(f'''{USER}/test-dynamic-config'''    ,      trust_remote_code=_lowerCamelCase			)
       # Can't make an isinstance check because the new_config is from the FakeConfig class of a dynamic module
       self.assertEqual(new_config.__class__.__name__    ,      '''CustomConfig'''			)
       self.assertEqual(new_config.attribute    ,      42			)






class       lowerCAmelCase__				(						unittest.TestCase	):







     def      lowercase	(							self			:			Optional[Any]			):
       _snake_case					     =       GPTaConfig()

       # attempt to modify each of int/float/bool/str config records and verify they were updated
       _snake_case					     =       c.n_embd + 1  # int
       _snake_case					     =       c.resid_pdrop + 1.0  # float
       _snake_case					     =       not c.scale_attn_weights  # bool
       _snake_case					     =       c.summary_type + '''foo'''  # str
       c.update_from_string(
           f'''n_embd={n_embd},resid_pdrop={resid_pdrop},scale_attn_weights={scale_attn_weights},summary_type={summary_type}'''			)
       self.assertEqual(_lowerCamelCase    ,      c.n_embd    ,      '''mismatch for key: n_embd'''			)
       self.assertEqual(_lowerCamelCase    ,      c.resid_pdrop    ,      '''mismatch for key: resid_pdrop'''			)
       self.assertEqual(_lowerCamelCase    ,      c.scale_attn_weights    ,      '''mismatch for key: scale_attn_weights'''			)
       self.assertEqual(_lowerCamelCase    ,      c.summary_type    ,      '''mismatch for key: summary_type'''			)







     def      lowercase	(							self			:			List[Any]			):
       _snake_case					     =       PretrainedConfig()
       _snake_case					     =       [key for key in base_config.__dict__ if key not in config_common_kwargs]
       # If this part of the test fails, you have arguments to addin config_common_kwargs above.
       self.assertListEqual(
           _lowerCamelCase    ,      ['''is_encoder_decoder''', '''_name_or_path''', '''_commit_hash''', '''transformers_version''']			)
       _snake_case					     =       [key for key, value in config_common_kwargs.items() if value == getattr(_lowerCamelCase    ,      _lowerCamelCase			)]
       if len(_lowerCamelCase			) > 0:
         raise ValueError(
             '''The following keys are set with the default values in'''
             ''' `test_configuration_common.config_common_kwargs` pick another value for them:'''
             f''' {', '.join(_lowerCamelCase			)}.'''			)







     def      lowercase	(							self			:			List[Any]			):
       with self.assertRaises(_lowerCamelCase			):
         # config is in subfolder, the following should not work without specifying the subfolder
         _snake_case					     =       BertConfig.from_pretrained('''hf-internal-testing/tiny-random-bert-subfolder'''			)

       _snake_case					     =       BertConfig.from_pretrained('''hf-internal-testing/tiny-random-bert-subfolder'''    ,      subfolder='''bert'''			)

       self.assertIsNotNone(_lowerCamelCase			)







     def      lowercase	(							self			:			Dict			):
       # A mock response for an HTTP head request to emulate server down
       _snake_case					     =       mock.Mock()
       _snake_case					     =       500
       _snake_case					     =       {}
       _snake_case					     =       HTTPError
       _snake_case					     =       {}

       # Download this model to make sure it's in the cache.
       _snake_case					     =       BertConfig.from_pretrained('''hf-internal-testing/tiny-random-bert'''			)

       # Under the mock environment we get a 500 error when trying to reach the model.
       with mock.patch('''requests.Session.request'''    ,      return_value=_lowerCamelCase			) as mock_head:
         _snake_case					     =       BertConfig.from_pretrained('''hf-internal-testing/tiny-random-bert'''			)
         # This check we did call the fake head request
         mock_head.assert_called()







     def      lowercase	(							self			:			List[Any]			):
       # This test is for deprecated behavior and can be removed in v5
       _snake_case					     =       BertConfig.from_pretrained(
           '''https://huggingface.co/hf-internal-testing/tiny-random-bert/resolve/main/config.json'''			)







     def      lowercase	(							self			:			Dict			):
       _snake_case					     =       AutoConfig.from_pretrained('''bert-base-cased'''			)
       _snake_case					     =       ['''config.4.0.0.json''']

       with tempfile.TemporaryDirectory() as tmp_dir:
         configuration.save_pretrained(_lowerCamelCase			)
         _snake_case					     =       2
         json.dump(configuration.to_dict()    ,      open(os.path.join(_lowerCamelCase    ,      '''config.4.0.0.json'''			)    ,      '''w'''			)			)

         # This should pick the new configuration file as the version of Transformers is > 4.0.0
         _snake_case					     =       AutoConfig.from_pretrained(_lowerCamelCase			)
         self.assertEqual(new_configuration.hidden_size    ,      2			)

         # Will need to be adjusted if we reach v42 and this test is still here.
         # Should pick the old configuration file as the version of Transformers is < 4.42.0
         _snake_case					     =       ['''config.42.0.0.json''']
         _snake_case					     =       768
         configuration.save_pretrained(_lowerCamelCase			)
         shutil.move(os.path.join(_lowerCamelCase    ,      '''config.4.0.0.json'''			)    ,      os.path.join(_lowerCamelCase    ,      '''config.42.0.0.json'''			)			)
         _snake_case					     =       AutoConfig.from_pretrained(_lowerCamelCase			)
         self.assertEqual(new_configuration.hidden_size    ,      768			)







     def      lowercase	(							self			:			Optional[int]			):
       # This repo has two configuration files, one for v4.0.0 and above with a different hidden size.
       _snake_case					     =       '''hf-internal-testing/test-two-configs'''

       import transformers as new_transformers

       _snake_case					     =       '''v4.0.0'''
       _snake_case     ,					_snake_case					     =       new_transformers.models.auto.AutoConfig.from_pretrained(
           _lowerCamelCase    ,      return_unused_kwargs=_lowerCamelCase			)
       self.assertEqual(new_configuration.hidden_size    ,      2			)
       # This checks `_configuration_file` ia not kept in the kwargs by mistake.
       self.assertDictEqual(_lowerCamelCase    ,      {}			)

       # Testing an older version by monkey-patching the version in the module it's used.
       import transformers as old_transformers

       _snake_case					     =       '''v3.0.0'''
       _snake_case					     =       old_transformers.models.auto.AutoConfig.from_pretrained(_lowerCamelCase			)
       self.assertEqual(old_configuration.hidden_size    ,      768			)





"""simple docstring"""



import argparse
import os
import re

import torch
from flax.traverse_util import flatten_dict
from tax import checkpoints

from transformers import (
    AutoTokenizer,
    PixaStructConfig,
    PixaStructForConditionalGeneration,
    PixaStructImageProcessor,
    PixaStructProcessor,
    PixaStructTextConfig,
    PixaStructVisionConfig,
)
def 					_UpperCAmelCase      ( __lowerCamelCase   :       Tuple		)     ->					Optional[int]:
  _snake_case					     =       checkpoints.load_tax_checkpoint(__lowerCamelCase		)
  _snake_case					     =       flatten_dict(__lowerCamelCase		)
  return flax_params
def 					_UpperCAmelCase      ( __lowerCamelCase   :       Dict		)     ->					Optional[int]:
  _snake_case					     =       {}

  _snake_case					     =       {
      '''token_embedder''': '''embeddings''',
      '''encoder_norm''': '''layernorm''',
      '''kernel''': '''weight''',
      '''.out''': '''.output''',
      '''scale''': '''weight''',
      '''embedders_0.pos_embedding''': '''row_embedder.weight''',
      '''embedders_1.pos_embedding''': '''column_embedder.weight''',
  }

  _snake_case					     =       {
      '''query''': '''attention.query''',
      '''key''': '''attention.key''',
      '''value''': '''attention.value''',
      '''output.dense''': '''output''',
      '''encoder_decoder_attention.o''': '''encoder_decoder_attention.attention.o''',
      '''pre_self_attention_layer_norm''': '''self_attention.layer_norm''',
      '''pre_cross_attention_layer_norm''': '''encoder_decoder_attention.layer_norm''',
      '''mlp.''': '''mlp.DenseReluDense.''',
      '''pre_mlp_layer_norm''': '''mlp.layer_norm''',
      '''self_attention.o''': '''self_attention.attention.o''',
      '''decoder.embeddings.embedding''': '''decoder.embed_tokens.weight''',
      '''decoder.relpos_bias.rel_embedding''': '''decoder.layer.0.self_attention.attention.relative_attention_bias.weight''',
      '''decoder.decoder_norm.weight''': '''decoder.final_layer_norm.weight''',
      '''decoder.logits_dense.weight''': '''decoder.lm_head.weight''',
  }

  for key in flax_dict.keys():
    if "target" in key:
      # remove the first prefix from the key
      _snake_case					     =       '''.'''.join(key[1:]		)

      # rename the key
      for old, new in CONVERSION_MAPPING.items():
        _snake_case					     =       new_key.replace(__lowerCamelCase ,					__lowerCamelCase		)

      if "decoder" in new_key:
        for old, new in DECODER_CONVERSION_MAPPING.items():
          _snake_case					     =       new_key.replace(__lowerCamelCase ,					__lowerCamelCase		)

      if "layers" in new_key and "decoder" not in new_key:
        # use regex to replace the layer number
        _snake_case					     =       re.sub(R'''layers_(\d+)''' ,					R'''layer.\1''' ,					__lowerCamelCase		)
        _snake_case					     =       new_key.replace('''encoder''' ,					'''encoder.encoder'''		)

      elif "layers" in new_key and "decoder" in new_key:
        # use regex to replace the layer number
        _snake_case					     =       re.sub(R'''layers_(\d+)''' ,					R'''layer.\1''' ,					__lowerCamelCase		)

      _snake_case					     =       flax_dict[key]

  _snake_case					     =       {}
  # convert converted_dict into torch format
  for key in converted_dict.keys():
    if ("embed_tokens" not in key) and ("embedder" not in key):
      _snake_case					     =       torch.from_numpy(converted_dict[key].T		)
    else:
      _snake_case					     =       torch.from_numpy(converted_dict[key]		)

  return converted_torch_dict
def 					_UpperCAmelCase      ( __lowerCamelCase   :       List[str] ,					__lowerCamelCase   :       Dict ,					__lowerCamelCase   :       Any=False ,					__lowerCamelCase   :       Optional[int]=False		)     ->					int:
  _snake_case					     =       get_flax_param(__lowerCamelCase		)

  if not use_large:
    _snake_case					     =       PixaStructVisionConfig()
    _snake_case					     =       PixaStructTextConfig()
  else:
    _snake_case					     =       PixaStructVisionConfig(
        hidden_size=15_36 ,					d_ff=39_68 ,					num_attention_heads=24 ,					num_hidden_layers=18		)
    _snake_case					     =       PixaStructTextConfig(hidden_size=15_36 ,					d_ff=39_68 ,					num_heads=24 ,					num_layers=18		)
  _snake_case					     =       PixaStructConfig(
      vision_config=encoder_config.to_dict() ,					text_config=decoder_config.to_dict() ,					is_vqa=__lowerCamelCase		)

  _snake_case					     =       PixaStructForConditionalGeneration(__lowerCamelCase		)

  _snake_case					     =       rename_and_convert_flax_params(__lowerCamelCase		)
  model.load_state_dict(__lowerCamelCase		)

  _snake_case					     =       AutoTokenizer.from_pretrained('''ybelkada/test-pix2struct-tokenizer'''		)
  _snake_case					     =       PixaStructImageProcessor()
  _snake_case					     =       PixaStructProcessor(image_processor=__lowerCamelCase ,					tokenizer=__lowerCamelCase		)

  if use_large:
    _snake_case					     =       40_96

  _snake_case					     =       True

  # mkdir if needed
  os.makedirs(__lowerCamelCase ,					exist_ok=__lowerCamelCase		)

  model.save_pretrained(__lowerCamelCase		)
  processor.save_pretrained(__lowerCamelCase		)

  print('''Model saved in {}'''.format(__lowerCamelCase		)		)


if __name__ == "__main__":
 UpperCAmelCase__          = argparse.ArgumentParser()
 parser.add_argument('--t5x_checkpoint_path', default=None, type=str, help='Path to the original T5x checkpoint.')
 parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
 parser.add_argument('--use_large', action='store_true', help='Use large model.')
 parser.add_argument('--is_vqa', action='store_true', help='Use large model.')
 UpperCAmelCase__          = parser.parse_args()

 convert_pixastruct_original_pytorch_checkpoint_to_hf(
     args.tax_checkpoint_path, args.pytorch_dump_folder_path, args.use_large
 )



