





import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaConfig, GPTaTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
    import jax
    import jax.numpy as jnp

    from transformers.modeling_flax_pytorch_utils import (
        convert_pytorch_state_dict_to_flax,
        load_flax_weights_in_pytorch_model,
    )
    from transformers.models.gpta.modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel

if is_torch_available():
    import torch







class 		_lowerCAmelCase:

  """simple docstring"""







  def __init__(   self    ,							UpperCAmelCase    ,							UpperCAmelCase=14    ,							UpperCAmelCase=7    ,							UpperCAmelCase=True    ,							UpperCAmelCase=True    ,							UpperCAmelCase=False    ,							UpperCAmelCase=True    ,							UpperCAmelCase=99    ,							UpperCAmelCase=32    ,							UpperCAmelCase=5    ,							UpperCAmelCase=4    ,							UpperCAmelCase=37    ,							UpperCAmelCase="gelu"    ,							UpperCAmelCase=0.1    ,							UpperCAmelCase=0.1    ,							UpperCAmelCase=5_12    ,							UpperCAmelCase=0.02    ,							)->   Dict:
     __A    =  parent
     __A    =  batch_size
     __A    =  seq_length
     __A    =  is_training
     __A    =  use_input_mask
     __A    =  use_token_type_ids
     __A    =  use_labels
     __A    =  vocab_size
     __A    =  hidden_size
     __A    =  num_hidden_layers
     __A    =  num_attention_heads
     __A    =  intermediate_size
     __A    =  hidden_act
     __A    =  hidden_dropout_prob
     __A    =  attention_probs_dropout_prob
     __A    =  max_position_embeddings
     __A    =  initializer_range
     __A    =  None
     __A    =  vocab_size - 1
     __A    =  vocab_size - 1
     __A    =  vocab_size - 1

  def 						SCREAMING_SNAKE_CASE__			(   self      )->   Dict:
     __A    =  ids_tensor([self.batch_size, self.seq_length]    ,							self.vocab_size      )

     __A    =  None
     if self.use_input_mask:
        __A    =  random_attention_mask([self.batch_size, self.seq_length]      )

     __A    =  GPTaConfig(
         vocab_size=self.vocab_size    ,							n_embd=self.hidden_size    ,							n_layer=self.num_hidden_layers    ,							n_head=self.num_attention_heads    ,							n_positions=self.max_position_embeddings    ,							use_cache=UpperCamelCase__    ,							bos_token_id=self.bos_token_id    ,							eos_token_id=self.eos_token_id    ,							pad_token_id=self.pad_token_id    ,							)

     return (config, input_ids, input_mask)

  def 						SCREAMING_SNAKE_CASE__			(   self      )->   Optional[Any]:
     __A    =  self.prepare_config_and_inputs()
     __A    =  config_and_inputs
     __A    =  {'input_ids': input_ids, 'attention_mask': attention_mask}
     return config, inputs_dict

  def 						SCREAMING_SNAKE_CASE__			(   self      )->   str:
     __A    =  self.prepare_config_and_inputs()

     __A    =  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
     __A    =  ids_tensor([self.batch_size, self.seq_length]    ,							vocab_size=2      )

     return (
         config,
         input_ids,
         attention_mask,
         encoder_hidden_states,
         encoder_attention_mask,
     )

  def 						SCREAMING_SNAKE_CASE__			(   self    ,							UpperCAmelCase    ,							UpperCAmelCase    ,							UpperCAmelCase    ,							UpperCAmelCase      )->   List[Any]:
     __A    =  20
     __A    =  model_class_name(UpperCamelCase__      )

     __A    =  model.init_cache(input_ids.shape[0]    ,							UpperCamelCase__      )
     __A    =  jnp.ones((input_ids.shape[0], max_decoder_length)    ,							dtype='''i4'''      )

     __A    =  jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1      )[None, :]    ,							(input_ids.shape[0], input_ids.shape[-1] - 1)      )
     __A    =  model(
         input_ids[:, :-1]    ,							attention_mask=UpperCamelCase__    ,							past_key_values=UpperCamelCase__    ,							position_ids=UpperCamelCase__    ,							)

     __A    =  jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]    ,							dtype='''i4'''      )
     __A    =  model(
         input_ids[:, -1:]    ,							attention_mask=UpperCamelCase__    ,							past_key_values=outputs_cache.past_key_values    ,							position_ids=UpperCamelCase__    ,							)

     __A    =  model(UpperCamelCase__      )

     __A    =  np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])      )      )
     self.parent.assertTrue(diff < 1e-3    ,							msg=f"Max diff is {diff}"      )





  def 						SCREAMING_SNAKE_CASE__			(   self    ,							UpperCAmelCase    ,							UpperCAmelCase    ,							UpperCAmelCase    ,							UpperCAmelCase      )->   Any:
     __A    =  20
     __A    =  model_class_name(UpperCamelCase__      )

     __A    =  jnp.concatenate(
         [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])      )]    ,							axis=-1    ,							)

     __A    =  model.init_cache(input_ids.shape[0]    ,							UpperCamelCase__      )
     __A    =  jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1      )[None, :]    ,							(input_ids.shape[0], input_ids.shape[-1] - 1)      )

     __A    =  model(
         input_ids[:, :-1]    ,							attention_mask=UpperCamelCase__    ,							past_key_values=UpperCamelCase__    ,							position_ids=UpperCamelCase__    ,							)
     __A    =  jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]    ,							dtype='''i4'''      )
     __A    =  model(
         input_ids[:, -1:]    ,							past_key_values=outputs_cache.past_key_values    ,							attention_mask=UpperCamelCase__    ,							position_ids=UpperCamelCase__    ,							)

     __A    =  model(UpperCamelCase__    ,							attention_mask=UpperCamelCase__      )

     __A    =  np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])      )      )
     self.parent.assertTrue(diff < 1e-3    ,							msg=f"Max diff is {diff}"      )







@require_flax
class 		_lowerCAmelCase(		__A						,       __A						,       unittest.TestCase):

  """simple docstring"""







  lowerCamelCase__     	=       (FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
  lowerCamelCase__     	=       (FlaxGPTaLMHeadModel,) if is_flax_available() else ()

  def 						SCREAMING_SNAKE_CASE__			(   self      )->   Any:
     __A    =  FlaxGPTaModelTester(self      )

  def 						SCREAMING_SNAKE_CASE__			(   self      )->   str:
     for model_class_name in self.all_model_classes:
        __A    =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_use_cache_forward(UpperCamelCase__    ,							UpperCamelCase__    ,							UpperCamelCase__    ,							UpperCamelCase__      )

  def 						SCREAMING_SNAKE_CASE__			(   self      )->   Any:
     for model_class_name in self.all_model_classes:
        __A    =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_use_cache_forward_with_attn_mask(
            UpperCamelCase__    ,							UpperCamelCase__    ,							UpperCamelCase__    ,							UpperCamelCase__      )

  @slow
  def 						SCREAMING_SNAKE_CASE__			(   self      )->   List[str]:
     __A    =  GPTaTokenizer.from_pretrained('''gpt2'''    ,							pad_token='''</s>'''    ,							padding_side='''left'''      )
     __A    =  tokenizer(['''Hello this is a long string''', '''Hey''']    ,							return_tensors='''np'''    ,							padding=UpperCamelCase__    ,							truncation=UpperCamelCase__      )

     __A    =  FlaxGPTaLMHeadModel.from_pretrained('''gpt2'''      )
     __A    =  False
     __A    =  model.config.eos_token_id

     __A    =  jax.jit(model.generate      )

     __A    =  jit_generate(inputs['''input_ids''']    ,							attention_mask=inputs['''attention_mask''']      ).sequences

     __A    =  tokenizer.batch_decode(UpperCamelCase__    ,							skip_special_tokens=UpperCamelCase__      )

     __A    =  [
         'Hello this is a long string of words. I\'m going to start with the first one.\n',
         'Hey, I\'m not sure if I\'m going to be able to do',
     ]

     self.assertListEqual(UpperCamelCase__    ,							UpperCamelCase__      )

  @is_pt_flax_cross_test
  def 						SCREAMING_SNAKE_CASE__			(   self      )->   Union[str, Any]:
     __A    =  self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__      ):
           # prepare inputs
           __A    =  self._prepare_for_class(UpperCamelCase__    ,							UpperCamelCase__      )
           __A    =  {k: torch.tensor(v.tolist()      ) for k, v in prepared_inputs_dict.items()}

           # load corresponding PyTorch class
           __A    =  model_class.__name__[4:]  # Skip the "Flax" at the beginning
           __A    =  getattr(UpperCamelCase__    ,							UpperCamelCase__      )

           __A    =  pt_inputs['input_ids'].shape
           __A    =  np.random.randint(0    ,							seq_length - 1    ,							size=(batch_size,)      )
           for batch_idx, start_index in enumerate(UpperCamelCase__      ):
              __A    =  0
              __A    =  1
              __A    =  0
              __A    =  1
           __A    =  pt_model_class(UpperCamelCase__      ).eval()
           __A    =  model_class(UpperCamelCase__    ,							dtype=jnp.floataa      )

           __A    =  convert_pytorch_state_dict_to_flax(pt_model.state_dict()    ,							UpperCamelCase__      )
           __A    =  fx_state

           with torch.no_grad():
              __A    =  pt_model(**UpperCamelCase__      ).to_tuple()

           __A    =  fx_model(**UpperCamelCase__      ).to_tuple()
           self.assertEqual(len(UpperCamelCase__      )    ,							len(UpperCamelCase__      )    ,							'''Output lengths differ between Flax and PyTorch'''      )
           for fx_output, pt_output in zip(UpperCamelCase__    ,							UpperCamelCase__      ):
              self.assert_almost_equals(fx_output[:, -1]    ,							pt_output[:, -1].numpy()    ,							4e-2      )

           with tempfile.TemporaryDirectory() as tmpdirname:
              pt_model.save_pretrained(UpperCamelCase__      )
              __A    =  model_class.from_pretrained(UpperCamelCase__    ,							from_pt=UpperCamelCase__      )

           __A    =  fx_model_loaded(**UpperCamelCase__      ).to_tuple()
           self.assertEqual(
               len(UpperCamelCase__      )    ,							len(UpperCamelCase__      )    ,							'''Output lengths differ between Flax and PyTorch'''      )
           for fx_output_loaded, pt_output in zip(UpperCamelCase__    ,							UpperCamelCase__      ):
              self.assert_almost_equals(fx_output_loaded[:, -1]    ,							pt_output[:, -1].numpy()    ,							4e-2      )

  @is_pt_flax_cross_test
  def 						SCREAMING_SNAKE_CASE__			(   self      )->   Optional[int]:
     __A    =  self.model_tester.prepare_config_and_inputs_for_common()
     for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__      ):
           # prepare inputs
           __A    =  self._prepare_for_class(UpperCamelCase__    ,							UpperCamelCase__      )
           __A    =  {k: torch.tensor(v.tolist()      ) for k, v in prepared_inputs_dict.items()}

           # load corresponding PyTorch class
           __A    =  model_class.__name__[4:]  # Skip the "Flax" at the beginning
           __A    =  getattr(UpperCamelCase__    ,							UpperCamelCase__      )

           __A    =  pt_model_class(UpperCamelCase__      ).eval()
           __A    =  model_class(UpperCamelCase__    ,							dtype=jnp.floataa      )

           __A    =  load_flax_weights_in_pytorch_model(UpperCamelCase__    ,							fx_model.params      )
           __A    =  pt_inputs['input_ids'].shape
           __A    =  np.random.randint(0    ,							seq_length - 1    ,							size=(batch_size,)      )
           for batch_idx, start_index in enumerate(UpperCamelCase__      ):
              __A    =  0
              __A    =  1
              __A    =  0
              __A    =  1

           # make sure weights are tied in PyTorch
           pt_model.tie_weights()

           with torch.no_grad():
              __A    =  pt_model(**UpperCamelCase__      ).to_tuple()

           __A    =  fx_model(**UpperCamelCase__      ).to_tuple()
           self.assertEqual(len(UpperCamelCase__      )    ,							len(UpperCamelCase__      )    ,							'''Output lengths differ between Flax and PyTorch'''      )
           for fx_output, pt_output in zip(UpperCamelCase__    ,							UpperCamelCase__      ):
              self.assert_almost_equals(fx_output[:, -1]    ,							pt_output[:, -1].numpy()    ,							4e-2      )

           with tempfile.TemporaryDirectory() as tmpdirname:
              fx_model.save_pretrained(UpperCamelCase__      )
              __A    =  pt_model_class.from_pretrained(UpperCamelCase__    ,							from_flax=UpperCamelCase__      )

           with torch.no_grad():
              __A    =  pt_model_loaded(**UpperCamelCase__      ).to_tuple()

           self.assertEqual(
               len(UpperCamelCase__      )    ,							len(UpperCamelCase__      )    ,							'''Output lengths differ between Flax and PyTorch'''      )
           for fx_output, pt_output in zip(UpperCamelCase__    ,							UpperCamelCase__      ):
              self.assert_almost_equals(fx_output[:, -1]    ,							pt_output[:, -1].numpy()    ,							4e-2      )





  @slow
  def 						SCREAMING_SNAKE_CASE__			(   self      )->   Tuple:
     for model_class_name in self.all_model_classes:
        __A    =  model_class_name.from_pretrained('''gpt2'''    ,							from_pt=UpperCamelCase__      )
        __A    =  model(np.ones((1, 1)      )      )
        self.assertIsNotNone(UpperCamelCase__      )









import argparse
import collections
import importlib.util
import os
import re


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_table.py
A										=      'src/diffusers'
A										=      'docs/source/en'
A										=      '.'



def        _lowerCamelCase(      lowerCAmelCase__					:		List[Any]     ,     lowerCAmelCase__					:		int     ,     lowerCAmelCase__					:		Tuple   ):





				'''simple docstring'''







				with open(lowerCAmelCase__     ,     'r'     ,     encoding='utf-8'     ,     newline='\n'   ) as f:
								SCREAMING_SNAKE_CASE_ :						Dict    				= f.readlines()
				# Find the start prompt.
				SCREAMING_SNAKE_CASE_ :						List[Any]    				= 0
				while not lines[start_index].startswith(lowerCAmelCase__   ):
								start_index += 1
				start_index += 1

				SCREAMING_SNAKE_CASE_ :						List[str]    				= start_index
				while not lines[end_index].startswith(lowerCAmelCase__   ):
								end_index += 1
				end_index -= 1

				while len(lines[start_index]   ) <= 1:
								start_index += 1
				while len(lines[end_index]   ) <= 1:
								end_index -= 1
				end_index += 1
				return "".join(lines[start_index:end_index]   ), start_index, end_index, lines


# Add here suffixes that are used to identify models, separated by |
A										=      'Model|Encoder|Decoder|ForConditionalGeneration'
# Regexes that match TF/Flax/PT model names.
A										=      re.compile(R'TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')
A										=      re.compile(R'Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')
# Will match any TF or Flax model too so need to be in an else branch afterthe two previous regexes.
A										=      re.compile(R'(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')


# This is to make sure the diffusers module imported is the one in the repo.
A										=      importlib.util.spec_from_file_location(
    'diffusers',
    os.path.join(TRANSFORMERS_PATH, '__init__.py'),
    submodule_search_locations=[TRANSFORMERS_PATH],
)
A										=      spec.loader.load_module()



def        _lowerCamelCase(      lowerCAmelCase__					:		Any   ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= re.finditer('.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)'     ,     lowerCAmelCase__   )
				return [m.group(0   ) for m in matches]



def        _lowerCamelCase(      lowerCAmelCase__					:		Optional[Any]     ,     lowerCAmelCase__					:		Tuple   ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_ :						Any    				= 2 if text == '✅' or text == '❌' else len(lowerCAmelCase__   )
				SCREAMING_SNAKE_CASE_ :						int    				= (width - text_length) // 2
				SCREAMING_SNAKE_CASE_ :						str    				= width - text_length - left_indent
				return " " * left_indent + text + " " * right_indent



def        _lowerCamelCase(      ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= diffusers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
				SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= {
				    name: config_mapping_names[code]
				    for code, name in diffusers_module.MODEL_NAMES_MAPPING.items()
				    if code in config_mapping_names
				}
				SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= {name: config.replace('ConfigMixin'     ,     ''   ) for name, config in model_name_to_config.items()}

				# Dictionaries flagging if each model prefix has a slow/fast tokenizer, backend in PT/TF/Flax.
				SCREAMING_SNAKE_CASE_ :						int    				= collections.defaultdict(lowerCAmelCase__   )
				SCREAMING_SNAKE_CASE_ :						List[Any]    				= collections.defaultdict(lowerCAmelCase__   )
				SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= collections.defaultdict(lowerCAmelCase__   )
				SCREAMING_SNAKE_CASE_ :						str    				= collections.defaultdict(lowerCAmelCase__   )
				SCREAMING_SNAKE_CASE_ :						Dict    				= collections.defaultdict(lowerCAmelCase__   )

				# Let's lookup through all diffusers object (once).
				for attr_name in dir(lowerCAmelCase__   ):
								SCREAMING_SNAKE_CASE_ :						List[str]    				= None
								if attr_name.endswith('Tokenizer'   ):
												SCREAMING_SNAKE_CASE_ :						List[Any]    				= slow_tokenizers
												SCREAMING_SNAKE_CASE_ :						Dict    				= attr_name[:-9]
								elif attr_name.endswith('TokenizerFast'   ):
												SCREAMING_SNAKE_CASE_ :						Any    				= fast_tokenizers
												SCREAMING_SNAKE_CASE_ :						int    				= attr_name[:-13]
								elif _re_tf_models.match(lowerCAmelCase__   ) is not None:
												SCREAMING_SNAKE_CASE_ :						List[Any]    				= tf_models
												SCREAMING_SNAKE_CASE_ :						Optional[int]    				= _re_tf_models.match(lowerCAmelCase__   ).groups()[0]
								elif _re_flax_models.match(lowerCAmelCase__   ) is not None:
												SCREAMING_SNAKE_CASE_ :						List[str]    				= flax_models
												SCREAMING_SNAKE_CASE_ :						int    				= _re_flax_models.match(lowerCAmelCase__   ).groups()[0]
								elif _re_pt_models.match(lowerCAmelCase__   ) is not None:
												SCREAMING_SNAKE_CASE_ :						Optional[int]    				= pt_models
												SCREAMING_SNAKE_CASE_ :						Any    				= _re_pt_models.match(lowerCAmelCase__   ).groups()[0]

								if lookup_dict is not None:
												while len(lowerCAmelCase__   ) > 0:
																if attr_name in model_name_to_prefix.values():
																				SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= True
																				break
																# Try again after removing the last word in the name
																SCREAMING_SNAKE_CASE_ :						str    				= ''.join(camel_case_split(lowerCAmelCase__   )[:-1]   )

    # Let's build that table!
				SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= list(model_name_to_config.keys()   )
				model_names.sort(key=str.lower   )
				SCREAMING_SNAKE_CASE_ :						List[Any]    				= ['Model', 'Tokenizer slow', 'Tokenizer fast', 'PyTorch support', 'TensorFlow support', 'Flax Support']
				# We'll need widths to properly display everything in the center (+2 is to leave one extra space on each side).
				SCREAMING_SNAKE_CASE_ :						str    				= [len(lowerCAmelCase__   ) + 2 for c in columns]
				SCREAMING_SNAKE_CASE_ :						List[Any]    				= max([len(lowerCAmelCase__   ) for name in model_names]   ) + 2

				# Build the table per se
				SCREAMING_SNAKE_CASE_ :						Dict    				= '|' + '|'.join([_center_text(lowerCAmelCase__     ,     lowerCAmelCase__   ) for c, w in zip(lowerCAmelCase__     ,     lowerCAmelCase__   )]   ) + '|\n'
				# Use ":-----:" format to center-aligned table cell texts
				table += "|" + "|".join([':' + '-' * (w - 2) + ':' for w in widths]   ) + "|\n"

				SCREAMING_SNAKE_CASE_ :						int    				= {True: '✅', False: '❌'}
				for name in model_names:
								SCREAMING_SNAKE_CASE_ :						int    				= model_name_to_prefix[name]
								SCREAMING_SNAKE_CASE_ :						int    				= [
								    name,
								    check[slow_tokenizers[prefix]],
								    check[fast_tokenizers[prefix]],
								    check[pt_models[prefix]],
								    check[tf_models[prefix]],
								    check[flax_models[prefix]],
								]
								table += "|" + "|".join([_center_text(lowerCAmelCase__     ,     lowerCAmelCase__   ) for l, w in zip(lowerCAmelCase__     ,     lowerCAmelCase__   )]   ) + "|\n"
				return table



def        _lowerCamelCase(      lowerCAmelCase__					:		Optional[int]=False   ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_ :						str    				= _find_text_in_file(
				    filename=os.path.join(lowerCAmelCase__     ,     'index.mdx'   )     ,     start_prompt='<!--This table is updated automatically from the auto modules'     ,     end_prompt='<!-- End table-->'     ,     )
				SCREAMING_SNAKE_CASE_ :						Tuple    				= get_model_table_from_auto_modules()

				if current_table != new_table:
								if overwrite:
												with open(os.path.join(lowerCAmelCase__     ,     'index.mdx'   )     ,     'w'     ,     encoding='utf-8'     ,     newline='\n'   ) as f:
																f.writelines(lines[:start_index] + [new_table] + lines[end_index:]   )
								else:
												raise ValueError(
												    'The model table in the `index.mdx` has not been updated. Run `make fix-copies` to fix this.'   )


if __name__ == "__main__":
				A										=      argparse.ArgumentParser()
				parser.add_argument('--fix_and_overwrite', action='store_true', help='Whether to fix inconsistencies.')
				A										=      parser.parse_args()

				check_model_table(args.fix_and_overwrite)