




'''simple docstring'''





import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
  import torch

if is_vision_available():
  from PIL import Image

  from transformers import ViltImageProcessor




class      UpperCamelCase       (	unittest.TestCase       ):



  def __init__(			self		,     __lowerCAmelCase		,     __lowerCAmelCase=7		,     __lowerCAmelCase=3		,     __lowerCAmelCase=18		,     __lowerCAmelCase=30		,     __lowerCAmelCase=400		,     __lowerCAmelCase=True		,     __lowerCAmelCase=None		,     __lowerCAmelCase=2		,     __lowerCAmelCase=True		,     __lowerCAmelCase=[0.5, 0.5, 0.5]		,     __lowerCAmelCase=[0.5, 0.5, 0.5]		,     ):







      '''simple docstring'''


      _a				:							Tuple							       =      size if size is not None else {"""shortest_edge""": 30}
      _a				:							Optional[int]							       =      parent
      _a				:							Union[str, Any]							       =      batch_size
      _a				:							int							       =      num_channels
      _a				:							List[str]							       =      image_size
      _a				:							str							       =      min_resolution
      _a				:							Union[str, Any]							       =      max_resolution
      _a				:							Tuple							       =      do_resize
      _a				:							str							       =      size
      _a				:							Any							       =      size_divisor
      _a				:							List[str]							       =      do_normalize
      _a				:							Tuple							       =      image_mean
      _a				:							str							       =      image_std



  def     a_							(			self ):







      '''simple docstring'''


      return {
          "image_mean": self.image_mean,
          "image_std": self.image_std,
          "do_normalize": self.do_normalize,
          "do_resize": self.do_resize,
          "size": self.size,
          "size_divisor": self.size_divisor,
      }



  def     a_							(			self		,     __lowerCAmelCase		,     __lowerCAmelCase=False ):







      '''simple docstring'''


      if not batched:
          _a				:							List[Any]							       =      self.size["""shortest_edge"""]
          _a				:							Optional[int]							       =      image_inputs[0]
          if isinstance(__lowerCAmelCase		,     Image.Image ):
              _a       , _a				:							Any							       =      image.size
          else:
              _a       , _a				:							Any							       =      image.shape[1], image.shape[2]
          _a				:							int							       =      size / min(__lowerCAmelCase		,     __lowerCAmelCase )
          if h < w:
              _a       , _a				:							str							       =      size, scale * w
          else:
              _a       , _a				:							Any							       =      scale * h, size

          _a				:							List[Any]							       =      int((1333 / 800) * size )
          if max(__lowerCAmelCase		,     __lowerCAmelCase ) > max_size:
              _a				:							Optional[Any]							       =      max_size / max(__lowerCAmelCase		,     __lowerCAmelCase )
              _a				:							Optional[int]							       =      newh * scale
              _a				:							List[str]							       =      neww * scale

          _a       , _a				:							int							       =      int(newh + 0.5 ), int(neww + 0.5 )
          _a       , _a				:							Dict							       =      (
              newh // self.size_divisor * self.size_divisor,
              neww // self.size_divisor * self.size_divisor,
          )

      else:
          _a				:							int							       =      []
          for image in image_inputs:
              _a       , _a				:							Any							       =      self.get_expected_values([image] )
              expected_values.append((expected_height, expected_width) )
          _a				:							Tuple							       =      max(__lowerCAmelCase		,     key=lambda __lowerCAmelCase : item[0] )[0]
          _a				:							List[str]							       =      max(__lowerCAmelCase		,     key=lambda __lowerCAmelCase : item[1] )[1]

      return expected_height, expected_width






@require_torch
@require_vision
class      UpperCamelCase       (	snake_case__      ,		unittest.TestCase       ):
  a     :Any          =					ViltImageProcessor if is_vision_available() else None



  def     a_							(			self ):







      '''simple docstring'''


      _a				:							Dict							       =      ViltImageProcessingTester(self )



  @property
  def     a_							(			self ):







      '''simple docstring'''


      return self.image_processor_tester.prepare_image_processor_dict()



  def     a_							(			self ):







      '''simple docstring'''


      _a				:							Tuple							       =      self.image_processing_class(**self.image_processor_dict )
      self.assertTrue(hasattr(__lowerCAmelCase		,     """image_mean""" ) )
      self.assertTrue(hasattr(__lowerCAmelCase		,     """image_std""" ) )
      self.assertTrue(hasattr(__lowerCAmelCase		,     """do_normalize""" ) )
      self.assertTrue(hasattr(__lowerCAmelCase		,     """do_resize""" ) )
      self.assertTrue(hasattr(__lowerCAmelCase		,     """size""" ) )
      self.assertTrue(hasattr(__lowerCAmelCase		,     """size_divisor""" ) )



  def     a_							(			self ):







      '''simple docstring'''


      _a				:							Any							       =      self.image_processing_class.from_dict(self.image_processor_dict )
      self.assertEqual(image_processor.size		,     {"""shortest_edge""": 30} )

      _a				:							List[Any]							       =      self.image_processing_class.from_dict(self.image_processor_dict		,     size=42 )
      self.assertEqual(image_processor.size		,     {"""shortest_edge""": 42} )



  def     a_							(			self ):







      '''simple docstring'''


      pass



  def     a_							(			self ):







      '''simple docstring'''


      _a				:							Optional[int]							       =      self.image_processing_class(**self.image_processor_dict )
      # create random PIL images
      _a				:							int							       =      prepare_image_inputs(self.image_processor_tester		,     equal_resolution=__lowerCAmelCase )
      for image in image_inputs:
          self.assertIsInstance(__lowerCAmelCase		,     Image.Image )

      # Test not batched input
      _a				:							str							       =      image_processing(image_inputs[0]		,     return_tensors="""pt""" ).pixel_values

      _a       , _a				:							int							       =      self.image_processor_tester.get_expected_values(__lowerCAmelCase )
      self.assertEqual(
          encoded_images.shape		,     (1, self.image_processor_tester.num_channels, expected_height, expected_width)		,     )

      # Test batched
      _a				:							Tuple							       =      image_processing(__lowerCAmelCase		,     return_tensors="""pt""" ).pixel_values

      _a       , _a				:							int							       =      self.image_processor_tester.get_expected_values(__lowerCAmelCase		,     batched=__lowerCAmelCase )
      self.assertEqual(
          encoded_images.shape		,     (
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              expected_height,
              expected_width,
          )		,     )



  def     a_							(			self ):







      '''simple docstring'''


      _a				:							Union[str, Any]							       =      self.image_processing_class(**self.image_processor_dict )
      # create random numpy tensors
      _a				:							Union[str, Any]							       =      prepare_image_inputs(self.image_processor_tester		,     equal_resolution=__lowerCAmelCase		,     numpify=__lowerCAmelCase )
      for image in image_inputs:
          self.assertIsInstance(__lowerCAmelCase		,     np.ndarray )

      # Test not batched input
      _a				:							int							       =      image_processing(image_inputs[0]		,     return_tensors="""pt""" ).pixel_values

      _a       , _a				:							List[Any]							       =      self.image_processor_tester.get_expected_values(__lowerCAmelCase )
      self.assertEqual(
          encoded_images.shape		,     (1, self.image_processor_tester.num_channels, expected_height, expected_width)		,     )

      # Test batched
      _a				:							List[Any]							       =      image_processing(__lowerCAmelCase		,     return_tensors="""pt""" ).pixel_values

      _a       , _a				:							List[str]							       =      self.image_processor_tester.get_expected_values(__lowerCAmelCase		,     batched=__lowerCAmelCase )
      self.assertEqual(
          encoded_images.shape		,     (
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              expected_height,
              expected_width,
          )		,     )



  def     a_							(			self ):







      '''simple docstring'''


      _a				:							str							       =      self.image_processing_class(**self.image_processor_dict )
      # create random PyTorch tensors
      _a				:							List[Any]							       =      prepare_image_inputs(self.image_processor_tester		,     equal_resolution=__lowerCAmelCase		,     torchify=__lowerCAmelCase )
      for image in image_inputs:
          self.assertIsInstance(__lowerCAmelCase		,     torch.Tensor )

      # Test not batched input
      _a				:							Optional[int]							       =      image_processing(image_inputs[0]		,     return_tensors="""pt""" ).pixel_values

      _a       , _a				:							Any							       =      self.image_processor_tester.get_expected_values(__lowerCAmelCase )
      self.assertEqual(
          encoded_images.shape		,     (1, self.image_processor_tester.num_channels, expected_height, expected_width)		,     )

      # Test batched
      _a				:							Optional[int]							       =      image_processing(__lowerCAmelCase		,     return_tensors="""pt""" ).pixel_values

      _a       , _a				:							List[Any]							       =      self.image_processor_tester.get_expected_values(__lowerCAmelCase		,     batched=__lowerCAmelCase )
      self.assertEqual(
          encoded_images.shape		,     (
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              expected_height,
              expected_width,
          )		,     )









'''simple docstring'''






import argparse
import os
from os.path import isfile

import torch

from transformers import GPTaConfig
def 				_lowerCamelCase(	UpperCamelCase__     :			List[str]     ,							UpperCamelCase__     :			Optional[Any]     ,							UpperCamelCase__     :			List[str]=0		)		->	Optional[int]:
    # Format the message.
    if name is None:
        A							:							Dict				=    None
    else:
        A							:							Optional[Any]				=    '''.''' * max(0     ,							spaces - 2		) + '''# {:''' + str(50 - spaces		) + '''s}'''
        A							:							Tuple				=    fmt.format(UpperCamelCase__		)

    # Print and recurse (if needed).
    if isinstance(UpperCamelCase__     ,							UpperCamelCase__		):
        if msg is not None:
            print(UpperCamelCase__		)
        for k in val.keys():
            recursive_print(UpperCamelCase__     ,							val[k]     ,							spaces + 2		)
    elif isinstance(UpperCamelCase__     ,							torch.Tensor		):
        print(UpperCamelCase__     ,							''':'''     ,							val.size()		)
    else:
        print(UpperCamelCase__     ,							''':'''     ,							UpperCamelCase__		)
def 				_lowerCamelCase(	UpperCamelCase__     :			Dict     ,							UpperCamelCase__     :			List[Any]     ,							UpperCamelCase__     :			Dict     ,							UpperCamelCase__     :			List[str]		)		->	Dict:
    # Permutes layout of param tensor to [num_splits * num_heads * hidden_size, :]
    # for compatibility with later versions of NVIDIA Megatron-LM.
    # The inverse operation is performed inside Megatron-LM to read checkpoints:
    # https://github.com/NVIDIA/Megatron-LM/blob/v2.4/megatron/checkpointing.py#L209
    # If param is the weight tensor of the self-attention block, the returned tensor
    # will have to be transposed one more time to be read by HuggingFace GPT2.
    A							:							int				=    param.size()
    # other versions store [num_heads * num_splits * hidden_size, :]
    A							:							List[str]				=    (num_heads, num_splits, hidden_size) + input_shape[1:]
    A							:							Optional[int]				=    param.view(*UpperCamelCase__		)
    A							:							Optional[Any]				=    param.transpose(0     ,							1		).contiguous()
    A							:							Dict				=    param.view(*UpperCamelCase__		)
    return param
def 				_lowerCamelCase(	UpperCamelCase__     :			Any     ,							UpperCamelCase__     :			Tuple		)		->	Optional[int]:
    A							:							int				=    config.n_positions
    A							:							Dict				=    config.n_layer
    A							:							List[str]				=    config.vocab_size
    A							:							Optional[Any]				=    config.n_head
    A							:							Union[str, Any]				=    config.n_embd // config.n_head

    A							:							Optional[int]				=    sd_megatron['''model.language_model.embedding.word_embeddings.weight'''][:vocab_size, :]
    A							:							Any				=    {
        '''transformer.wte.weight''': word_embeddings,
        '''transformer.wpe.weight''': sd_megatron['''model.language_model.embedding.position_embeddings.weight'''],
        '''transformer.ln_f.weight''': sd_megatron['''model.language_model.encoder.final_layernorm.weight'''],
        '''transformer.ln_f.bias''': sd_megatron['''model.language_model.encoder.final_layernorm.bias'''],
    }

    A							:							Optional[int]				=    '''model.language_model.encoder.layers.'''
    for i in range(UpperCamelCase__		):
        A							:							Dict				=    torch.tril(torch.ones((n_positions, n_positions)     ,							dtype=torch.bool		)		)
        A							:							Any				=    causal_mask.view(1     ,							1     ,							UpperCamelCase__     ,							UpperCamelCase__		)
        A							:							List[Any]				=    causal_mask
        A							:							List[str]				=    torch.tensor(-1e4     ,							dtype=torch.bfloataa		)

        A							:							Optional[int]				=    sd_megatron[F'''{pf}{i}.input_layernorm.weight''']
        A							:							Optional[int]				=    sd_megatron[F'''{pf}{i}.input_layernorm.bias''']

        A							:							Any				=    sd_megatron[F'''{pf}{i}.self_attention.query_key_value.weight''']
        A							:							List[Any]				=    fix_query_key_value_ordering(UpperCamelCase__     ,							3     ,							UpperCamelCase__     ,							UpperCamelCase__		)
        A							:							List[Any]				=    vala.transpose(0     ,							1		).contiguous()

        A							:							List[Any]				=    sd_megatron[F'''{pf}{i}.self_attention.query_key_value.bias''']
        A							:							Any				=    fix_query_key_value_ordering(UpperCamelCase__     ,							3     ,							UpperCamelCase__     ,							UpperCamelCase__		)
        A							:							List[Any]				=    vala

        A							:							Tuple				=    sd_megatron[F'''{pf}{i}.self_attention.dense.weight'''].transpose(
            0     ,							1		)
        A							:							Optional[int]				=    sd_megatron[F'''{pf}{i}.self_attention.dense.bias''']
        A							:							Any				=    sd_megatron[F'''{pf}{i}.post_attention_layernorm.weight''']
        A							:							Any				=    sd_megatron[F'''{pf}{i}.post_attention_layernorm.bias''']
        A							:							str				=    sd_megatron[F'''{pf}{i}.mlp.dense_h_to_4h.weight'''].transpose(0     ,							1		)
        A							:							List[Any]				=    sd_megatron[F'''{pf}{i}.mlp.dense_h_to_4h.bias''']
        A							:							Tuple				=    sd_megatron[F'''{pf}{i}.mlp.dense_4h_to_h.weight'''].transpose(
            0     ,							1		)
        A							:							Union[str, Any]				=    sd_megatron[F'''{pf}{i}.mlp.dense_4h_to_h.bias''']

    # For LM head, transformers' wants the matrix to weight embeddings.
    A							:							Optional[int]				=    word_embeddings

    return sd_hf
def 				_lowerCamelCase(	UpperCamelCase__     :			Union[str, Any]     ,							UpperCamelCase__     :			Optional[Any]		)		->	Union[str, Any]:
    A							:							List[str]				=    64_000
    A							:							str				=    config_megatron['''encoder_seq_length''']
    A							:							Union[str, Any]				=    config_megatron['''hidden_size''']
    A							:							Tuple				=    config_megatron['''num_layers''']
    A							:							List[str]				=    config_megatron['''num_attention_heads''']
    A							:							List[Any]				=    config_megatron['''ffn_hidden_size''']
    A							:							List[str]				=    '''gelu'''
    A							:							int				=    0.1
    A							:							Union[str, Any]				=    0.1
    A							:							Optional[int]				=    0.1
    A							:							Optional[int]				=    config_megatron['''layernorm_epsilon''']  # 1e-5
    A							:							Tuple				=    config_megatron['''init_method_std''']  # 0.02
    A							:							Any				=    config_megatron['''apply_query_key_layer_scaling''']  # True
    A							:							Tuple				=    True
    A							:							str				=    True

    # This identifies the 6.7B (7B) model which uses a different tokenizer
    if config_megatron["hidden_size"] == 4_096:
        A							:							Tuple				=    1  # <|endoftext|>
        A							:							Any				=    1  # <|endoftext|>
        A							:							List[str]				=    0  # <unk>
    else:
        A							:							Any				=    2  # <s>
        A							:							str				=    3  # <|endoftext|>
        A							:							Tuple				=    0  # <pad>

    return config_hf
def 				_lowerCamelCase(	UpperCamelCase__     :			List[Any]		)		->	Tuple:
    print(UpperCamelCase__		)

    A							:							Tuple				=    args.checkpoint_path
    A							:							Dict				=    args.save_path
    if isfile(UpperCamelCase__		):
        raise FileNotFoundError(F'''ERROR! could not find file {checkpoint_path}'''		)

    # Load the model.
    A							:							str				=    torch.load(UpperCamelCase__     ,							map_location='''cpu'''		)

    # Load the config.
    A							:							List[str]				=    checkpoint['''hyper_parameters''']['''cfg''']
    A							:							Tuple				=    GPTaConfig()
    A							:							int				=    copy_config(config_hf=UpperCamelCase__     ,							config_megatron=UpperCamelCase__		)
    A							:							Any				=    ['''GPT2LMHeadModel''']

    A							:							Optional[int]				=    checkpoint['''state_dict''']

    # Convert.
    print('''Converting'''		)
    A							:							Any				=    convert_megatron_checkpoint(UpperCamelCase__     ,							UpperCamelCase__		)

    # Print the structure of converted state dict.
    if args.print_checkpoint_structure:
        recursive_print(UpperCamelCase__     ,							UpperCamelCase__		)

    A							:							Union[str, Any]				=    '''GPTSw3Tokenizer'''

    # Store the config to file.
    print('''Saving config'''		)
    config_hf.save_pretrained(UpperCamelCase__		)

    # Store the state_dict to file.
    A							:							int				=    os.path.join(UpperCamelCase__     ,							'''pytorch_model.bin'''		)
    print(F'''Saving checkpoint to "{output_checkpoint_file}"'''		)
    torch.save(UpperCamelCase__     ,							UpperCamelCase__		)


if __name__ == "__main__":
     snake_case_     			=						argparse.ArgumentParser()
     parser.add_argument(
         """--checkpoint_path""",
         type=str,
         required=True,
         help="""e.g. megatron_gpt--val_loss=2.42-step=38000-consumed_samples=54720000""",
     )
     parser.add_argument("""--save_path""", type=str, required=True, help="""e.g. /home/user/gpt-sw3/hf""")
     parser.add_argument("""--print-checkpoint-structure""", action="""store_true""")
     snake_case_     			=						parser.parse_args()
     main(_args)
