





"""simple docstring"""




import collections
import inspect
import unittest

from transformers import DonutSwinConfig
from transformers.testing_utils import require_torch, slow, torch_device
from transformers.utils import is_torch_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch
   from torch import nn

   from transformers import DonutSwinModel
   from transformers.models.donut.modeling_donut_swin import DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST

class 					_UpperCAmelCase    :



      '''simple docstring'''

      def __init__(	self   , snake_case_   , snake_case_=1_3   , snake_case_=3_2   , snake_case_=2   , snake_case_=3   , snake_case_=1_6   , snake_case_=[1, 2, 1]   , snake_case_=[2, 2, 4]   , snake_case_=2   , snake_case_=2.0   , snake_case_=True   , snake_case_=0.0   , snake_case_=0.0   , snake_case_=0.1   , snake_case_="gelu"   , snake_case_=False   , snake_case_=True   , snake_case_=0.02   , snake_case_=1E-5   , snake_case_=True   , snake_case_=None   , snake_case_=True   , snake_case_=1_0   , snake_case_=8   , ):

        """simple docstring"""

        A_    :			str									=    parent
        A_    :			Optional[int]									=    batch_size
        A_    :			Optional[int]									=    image_size
        A_    :			Optional[Any]									=    patch_size
        A_    :			List[str]									=    num_channels
        A_    :			List[str]									=    embed_dim
        A_    :			int									=    depths
        A_    :			Tuple									=    num_heads
        A_    :			int									=    window_size
        A_    :			Tuple									=    mlp_ratio
        A_    :			str									=    qkv_bias
        A_    :			List[str]									=    hidden_dropout_prob
        A_    :			List[str]									=    attention_probs_dropout_prob
        A_    :			Tuple									=    drop_path_rate
        A_    :			Dict									=    hidden_act
        A_    :			Optional[int]									=    use_absolute_embeddings
        A_    :			Union[str, Any]									=    patch_norm
        A_    :			Union[str, Any]									=    layer_norm_eps
        A_    :			Tuple									=    initializer_range
        A_    :			int									=    is_training
        A_    :			List[Any]									=    scope
        A_    :			Optional[Any]									=    use_labels
        A_    :			str									=    type_sequence_label_size
        A_    :			List[Any]									=    encoder_stride

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Tuple									=    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

        A_    :			Tuple									=    None
        if self.use_labels:
          A_    :			List[Any]									=    ids_tensor([self.batch_size]   , self.type_sequence_label_size    )

        A_    :			Optional[Any]									=    self.get_config()

        return config, pixel_values, labels

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        return DonutSwinConfig(
            image_size=self.image_size   , patch_size=self.patch_size   , num_channels=self.num_channels   , embed_dim=self.embed_dim   , depths=self.depths   , num_heads=self.num_heads   , window_size=self.window_size   , mlp_ratio=self.mlp_ratio   , qkv_bias=self.qkv_bias   , hidden_dropout_prob=self.hidden_dropout_prob   , attention_probs_dropout_prob=self.attention_probs_dropout_prob   , drop_path_rate=self.drop_path_rate   , hidden_act=self.hidden_act   , use_absolute_embeddings=self.use_absolute_embeddings   , path_norm=self.patch_norm   , layer_norm_eps=self.layer_norm_eps   , initializer_range=self.initializer_range   , encoder_stride=self.encoder_stride   , )

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_   , snake_case_    ):

        """simple docstring"""

        A_    :			Any									=    DonutSwinModel(config=__UpperCamelCase    )
        model.to(__UpperCamelCase    )
        model.eval()
        A_    :			List[str]									=    model(__UpperCamelCase    )

        A_    :			List[str]									=    ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths    ) - 1))
        A_    :			Tuple									=    int(config.embed_dim * 2 ** (len(config.depths    ) - 1)    )

        self.parent.assertEqual(result.last_hidden_state.shape   , (self.batch_size, expected_seq_len, expected_dim)    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			int									=    self.prepare_config_and_inputs()
        (
            (
            A_	
        )		,			(
            A_	
        )		,			(
            A_	
        )		,				
        )    :			Tuple									=    config_and_inputs
        A_    :			Tuple									=    {'pixel_values': pixel_values}
        return config, inputs_dict


@require_torch
class 					_UpperCAmelCase    (  __snake_case     ,  __snake_case     ,  unittest.TestCase						):



      '''simple docstring'''
      lowercase_ :					int             =  (DonutSwinModel,) if is_torch_available() else ()
      lowercase_ :					Any             =  {"""feature-extraction""": DonutSwinModel} if is_torch_available() else {}
      lowercase_ :					str             =  True

      lowercase_ :					Optional[Any]             =  False
      lowercase_ :					int             =  False
      lowercase_ :					Dict             =  False

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Optional[Any]									=    DonutSwinModelTester(self    )
        A_    :			str									=    ConfigTester(self   , config_class=__UpperCamelCase   , embed_dim=3_7    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        self.create_and_test_config_common_properties()
        self.config_tester.create_and_test_config_to_json_string()
        self.config_tester.create_and_test_config_to_json_file()
        self.config_tester.create_and_test_config_from_and_save_pretrained()
        self.config_tester.create_and_test_config_with_num_labels()
        self.config_tester.check_config_can_be_init_without_params()
        self.config_tester.check_config_arguments_init()

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        return

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			List[str]									=    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__UpperCamelCase    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        pass

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_		,			A_    :			Union[str, Any]									=    self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          A_    :			Dict									=    model_class(__UpperCamelCase    )
          self.assertIsInstance(model.get_input_embeddings()   , (nn.Module)    )
          A_    :			Any									=    model.get_output_embeddings()
          self.assertTrue(x is None or isinstance(__UpperCamelCase   , nn.Linear    )    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_		,			A_    :			List[str]									=    self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          A_    :			Tuple									=    model_class(__UpperCamelCase    )
          A_    :			Dict									=    inspect.signature(model.forward    )
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          A_    :			Tuple									=    [*signature.parameters.keys()]

          A_    :			List[Any]									=    ['pixel_values']
          self.assertListEqual(arg_names[:1]   , __UpperCamelCase    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_		,			A_    :			int									=    self.model_tester.prepare_config_and_inputs_for_common()
        A_    :			Tuple									=    True

        for model_class in self.all_model_classes:
          A_    :			Optional[int]									=    True
          A_    :			Union[str, Any]									=    False
          A_    :			Union[str, Any]									=    True
          A_    :			Optional[int]									=    model_class(__UpperCamelCase    )
          model.to(__UpperCamelCase    )
          model.eval()
          with torch.no_grad():
            A_    :			str									=    model(**self._prepare_for_class(__UpperCamelCase   , __UpperCamelCase    )    )
          A_    :			List[Any]									=    outputs.attentions
          A_    :			int									=    len(self.model_tester.depths    )
          self.assertEqual(len(__UpperCamelCase    )   , __UpperCamelCase    )

          # check that output_attentions also work using config
          del inputs_dict["output_attentions"]
          A_    :			str									=    True
          A_    :			Tuple									=    config.window_size**2
          A_    :			Optional[int]									=    model_class(__UpperCamelCase    )
          model.to(__UpperCamelCase    )
          model.eval()
          with torch.no_grad():
            A_    :			Optional[Any]									=    model(**self._prepare_for_class(__UpperCamelCase   , __UpperCamelCase    )    )
          A_    :			Dict									=    outputs.attentions
          self.assertEqual(len(__UpperCamelCase    )   , __UpperCamelCase    )

          self.assertListEqual(
              list(attentions[0].shape[-3:]    )   , [self.model_tester.num_heads[0], window_size_squared, window_size_squared]   , )
          A_    :			Union[str, Any]									=    len(__UpperCamelCase    )

          # Check attention is always last and order is fine
          A_    :			Union[str, Any]									=    True
          A_    :			Dict									=    True
          A_    :			List[str]									=    model_class(__UpperCamelCase    )
          model.to(__UpperCamelCase    )
          model.eval()
          with torch.no_grad():
            A_    :			Optional[Any]									=    model(**self._prepare_for_class(__UpperCamelCase   , __UpperCamelCase    )    )

          if hasattr(self.model_tester   , 'num_hidden_states_types'    ):
            A_    :			Tuple									=    self.model_tester.num_hidden_states_types
          else:
            # also another +1 for reshaped_hidden_states
            A_    :			Optional[int]									=    2
          self.assertEqual(out_len + added_hidden_states   , len(__UpperCamelCase    )    )

          A_    :			List[Any]									=    outputs.attentions

          self.assertEqual(len(__UpperCamelCase    )   , __UpperCamelCase    )

          self.assertListEqual(
              list(self_attentions[0].shape[-3:]    )   , [self.model_tester.num_heads[0], window_size_squared, window_size_squared]   , )

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_   , snake_case_   , snake_case_    ):

        """simple docstring"""

        A_    :			Tuple									=    model_class(__UpperCamelCase    )
        model.to(__UpperCamelCase    )
        model.eval()

        with torch.no_grad():
          A_    :			str									=    model(**self._prepare_for_class(__UpperCamelCase   , __UpperCamelCase    )    )

        A_    :			Dict									=    outputs.hidden_states

        A_    :			str									=    getattr(
            self.model_tester   , 'expected_num_hidden_layers'   , len(self.model_tester.depths    ) + 1    )
        self.assertEqual(len(__UpperCamelCase    )   , __UpperCamelCase    )

        # DonutSwin has a different seq_length
        A_    :			int									=    (
            config.patch_size
            if isinstance(config.patch_size   , collections.abc.Iterable    )
            else (config.patch_size, config.patch_size)
        )

        A_    :			Tuple									=    (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

        self.assertListEqual(
            list(hidden_states[0].shape[-2:]    )   , [num_patches, self.model_tester.embed_dim]   , )

        A_    :			int									=    outputs.reshaped_hidden_states
        self.assertEqual(len(__UpperCamelCase    )   , __UpperCamelCase    )

        A_		,			A_		,			A_		,			A_    :			Optional[int]									=    reshaped_hidden_states[0].shape
        A_    :			List[str]									=    (
            reshaped_hidden_states[0].view(__UpperCamelCase   , __UpperCamelCase   , height * width    ).permute(0   , 2   , 1    )
        )
        self.assertListEqual(
            list(reshaped_hidden_states.shape[-2:]    )   , [num_patches, self.model_tester.embed_dim]   , )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_		,			A_    :			Optional[int]									=    self.model_tester.prepare_config_and_inputs_for_common()

        A_    :			Optional[int]									=    (
            self.model_tester.image_size
            if isinstance(self.model_tester.image_size   , collections.abc.Iterable    )
            else (self.model_tester.image_size, self.model_tester.image_size)
        )

        for model_class in self.all_model_classes:
          A_    :			List[str]									=    True
          self.check_hidden_states_output(__UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase    )

          # check that output_hidden_states also work using config
          del inputs_dict["output_hidden_states"]
          A_    :			List[Any]									=    True

          self.check_hidden_states_output(__UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_		,			A_    :			Tuple									=    self.model_tester.prepare_config_and_inputs_for_common()
        A_    :			Tuple									=    3

        A_    :			Any									=    (
            self.model_tester.image_size
            if isinstance(self.model_tester.image_size   , collections.abc.Iterable    )
            else (self.model_tester.image_size, self.model_tester.image_size)
        )
        A_    :			Tuple									=    (
            config.patch_size
            if isinstance(config.patch_size   , collections.abc.Iterable    )
            else (config.patch_size, config.patch_size)
        )

        A_    :			Any									=    image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
        A_    :			Optional[int]									=    image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

        for model_class in self.all_model_classes:
          A_    :			Tuple									=    True
          self.check_hidden_states_output(__UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase   , (padded_height, padded_width)    )

          # check that output_hidden_states also work using config
          del inputs_dict["output_hidden_states"]
          A_    :			List[str]									=    True
          self.check_hidden_states_output(__UpperCamelCase   , __UpperCamelCase   , __UpperCamelCase   , (padded_height, padded_width)    )

      @slow
      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        for model_name in DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          A_    :			Optional[Any]									=    DonutSwinModel.from_pretrained(__UpperCamelCase    )
          self.assertIsNotNone(__UpperCamelCase    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_		,			A_    :			Optional[Any]									=    self.model_tester.prepare_config_and_inputs_for_common()

        A_    :			List[str]									=    _config_zero_init(__UpperCamelCase    )
        for model_class in self.all_model_classes:
          A_    :			str									=    model_class(config=__UpperCamelCase    )
          for name, param in model.named_parameters():
            if "embeddings" not in name and param.requires_grad:
              self.assertIn(
                  ((param.data.mean() * 1E9).round() / 1E9).item()   , [0.0, 1.0]   , msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""   , )

import argparse
import os
from os.path import isfile

import torch

from transformers import GPTaConfig





def        lowerCAmelCase       (				snake_case__   :      List[str]				,					snake_case__   :      List[str]				,					snake_case__   :      Optional[int]=0 )->			Union[str, Any]:
    # Format the message.
    if name is None:
        A_								=					None
    else:
        A_								=					"." * max(0				,					spaces - 2 ) + "# {:" + str(50 - spaces ) + "s}"
        A_								=					fmt.format(snake_case__ )

    # Print and recurse (if needed).
    if isinstance(snake_case__				,					snake_case__ ):
        if msg is not None:
            print(snake_case__ )
        for k in val.keys():
            recursive_print(snake_case__				,					val[k]				,					spaces + 2 )
    elif isinstance(snake_case__				,					torch.Tensor ):
        print(snake_case__				,					":"				,					val.size() )
    else:
        print(snake_case__				,					":"				,					snake_case__ )





def        lowerCAmelCase       (				snake_case__   :      Any				,					snake_case__   :      List[Any]				,					snake_case__   :      str				,					snake_case__   :      str )->			Union[str, Any]:
    # 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_								=					param.size()
    # other versions store [num_heads * num_splits * hidden_size, :]
    A_								=					(num_heads, num_splits, hidden_size) + input_shape[1:]
    A_								=					param.view(*snake_case__ )
    A_								=					param.transpose(0				,					1 ).contiguous()
    A_								=					param.view(*snake_case__ )
    return param





def        lowerCAmelCase       (				snake_case__   :      Dict				,					snake_case__   :      str )->			Optional[int]:
    A_								=					config.n_positions
    A_								=					config.n_layer
    A_								=					config.vocab_size
    A_								=					config.n_head
    A_								=					config.n_embd // config.n_head

    A_								=					sd_megatron["model.language_model.embedding.word_embeddings.weight"][:vocab_size, :]
    A_								=					{
        "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_								=					"model.language_model.encoder.layers."
    for i in range(snake_case__ ):
        A_								=					torch.tril(torch.ones((n_positions, n_positions)				,					dtype=torch.bool ) )
        A_								=					causal_mask.view(1				,					1				,					snake_case__				,					snake_case__ )
        A_								=					causal_mask
        A_								=					torch.tensor(-1e4				,					dtype=torch.bfloataa )

        A_								=					sd_megatron[f'{pf}{i}.input_layernorm.weight']
        A_								=					sd_megatron[f'{pf}{i}.input_layernorm.bias']

        A_								=					sd_megatron[f'{pf}{i}.self_attention.query_key_value.weight']
        A_								=					fix_query_key_value_ordering(snake_case__				,					3				,					snake_case__				,					snake_case__ )
        A_								=					vala.transpose(0				,					1 ).contiguous()

        A_								=					sd_megatron[f'{pf}{i}.self_attention.query_key_value.bias']
        A_								=					fix_query_key_value_ordering(snake_case__				,					3				,					snake_case__				,					snake_case__ )
        A_								=					vala

        A_								=					sd_megatron[f'{pf}{i}.self_attention.dense.weight'].transpose(
            0				,					1 )
        A_								=					sd_megatron[f'{pf}{i}.self_attention.dense.bias']
        A_								=					sd_megatron[f'{pf}{i}.post_attention_layernorm.weight']
        A_								=					sd_megatron[f'{pf}{i}.post_attention_layernorm.bias']
        A_								=					sd_megatron[f'{pf}{i}.mlp.dense_h_to_4h.weight'].transpose(0				,					1 )
        A_								=					sd_megatron[f'{pf}{i}.mlp.dense_h_to_4h.bias']
        A_								=					sd_megatron[f'{pf}{i}.mlp.dense_4h_to_h.weight'].transpose(
            0				,					1 )
        A_								=					sd_megatron[f'{pf}{i}.mlp.dense_4h_to_h.bias']

    # For LM head, transformers' wants the matrix to weight embeddings.
    A_								=					word_embeddings

    return sd_hf





def        lowerCAmelCase       (				snake_case__   :      Optional[int]				,					snake_case__   :      List[Any] )->			Dict:
    A_								=					64000
    A_								=					config_megatron["encoder_seq_length"]
    A_								=					config_megatron["hidden_size"]
    A_								=					config_megatron["num_layers"]
    A_								=					config_megatron["num_attention_heads"]
    A_								=					config_megatron["ffn_hidden_size"]
    A_								=					"gelu"
    A_								=					0.1
    A_								=					0.1
    A_								=					0.1
    A_								=					config_megatron["layernorm_epsilon"]  # 1e-5
    A_								=					config_megatron["init_method_std"]  # 0.02
    A_								=					config_megatron["apply_query_key_layer_scaling"]  # True
    A_								=					True
    A_								=					True

    # This identifies the 6.7B (7B) model which uses a different tokenizer
    if config_megatron["hidden_size"] == 4096:
        A_								=					1  # <|endoftext|>
        A_								=					1  # <|endoftext|>
        A_								=					0  # <unk>
    else:
        A_								=					2  # <s>
        A_								=					3  # <|endoftext|>
        A_								=					0  # <pad>

    return config_hf





def        lowerCAmelCase       (				snake_case__   :      Tuple )->			Dict:
    print(snake_case__ )

    A_								=					args.checkpoint_path
    A_								=					args.save_path
    if isfile(snake_case__ ):
        raise FileNotFoundError(f'ERROR! could not find file {checkpoint_path}' )

    # Load the model.
    A_								=					torch.load(snake_case__				,					map_location="cpu" )

    # Load the config.
    A_								=					checkpoint["hyper_parameters"]["cfg"]
    A_								=					GPTaConfig()
    A_								=					copy_config(config_hf=snake_case__				,					config_megatron=snake_case__ )
    A_								=					["GPT2LMHeadModel"]

    A_								=					checkpoint["state_dict"]

    # Convert.
    print("Converting" )
    A_								=					convert_megatron_checkpoint(snake_case__				,					snake_case__ )

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

    A_								=					"GPTSw3Tokenizer"

    # Store the config to file.
    print("Saving config" )
    config_hf.save_pretrained(snake_case__ )

    # Store the state_dict to file.
    A_								=					os.path.join(snake_case__				,					"pytorch_model.bin" )
    print(f'Saving checkpoint to "{output_checkpoint_file}"' )
    torch.save(snake_case__				,					snake_case__ )


if __name__ == "__main__":
     __magic_name__  :  Union[str, Any]     			=   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')
     __magic_name__  :  str     			=   parser.parse_args()
     main(_args)






