






"""simple docstring"""






import warnings
from typing import List
from unittest.mock import Mock

import torch
from torch.utils.data import DataLoader, IterableDataset, TensorDataset

from accelerate.accelerator import Accelerator
from accelerate.utils.dataclasses import DistributedType




class 							lowerCAmelCase__				(			__snake_case ):






      '''simple docstring'''

      def __init__(						self     :  Dict						,	lowercase_     :  str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		data



      def __iter__(						self     :  Dict):




          '''simple docstring'''


          for element in self.data:
              yield element



def 					_A			(__a=True							) ->       List[Any]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  Dict    						=		Accelerator(even_batches=__SCREAMING_SNAKE_CASE							)
    assert accelerator.num_processes == 2, "this script expects that two GPUs are available"
    return accelerator



def 					_A			(__a    ,				__a    ,				__a    ,				__a = False							) ->       Dict:
    """simple docstring"""






    if iterable:
        SCREAMING_SNAKE_CASE_      :  Any    						=		DummyIterableDataset(torch.as_tensor(range(__SCREAMING_SNAKE_CASE							)							)							)
    else:
        SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		TensorDataset(torch.as_tensor(range(__SCREAMING_SNAKE_CASE							)							)							)

    SCREAMING_SNAKE_CASE_      :  Any    						=		DataLoader(__SCREAMING_SNAKE_CASE    ,				batch_size=__SCREAMING_SNAKE_CASE							)
    SCREAMING_SNAKE_CASE_      :  List[str]    						=		accelerator.prepare(__SCREAMING_SNAKE_CASE							)

    return dl



def 					_A			(__a    ,				__a    ,				__a    ,				__a    ,				__a    ,				) ->       int:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  Tuple    						=		create_dataloader(accelerator=__SCREAMING_SNAKE_CASE    ,				dataset_size=__SCREAMING_SNAKE_CASE    ,				batch_size=__SCREAMING_SNAKE_CASE							)

    SCREAMING_SNAKE_CASE_      :  str    						=		[len(batch[0]							) for batch in dl]

    if accelerator.process_index == 0:
        assert batch_sizes == process_0_expected_batch_sizes
    elif accelerator.process_index == 1:
        assert batch_sizes == process_1_expected_batch_sizes



def 					_A			() ->       Optional[Any]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  str    						=		create_accelerator()

    # without padding, we would expect a different number of batches
    verify_dataloader_batch_sizes(
        __SCREAMING_SNAKE_CASE    ,				dataset_size=3    ,				batch_size=1    ,				process_0_expected_batch_sizes=[1, 1]    ,				process_1_expected_batch_sizes=[1, 1]    ,				)

    # without padding, we would expect the same number of batches, but different sizes
    verify_dataloader_batch_sizes(
        __SCREAMING_SNAKE_CASE    ,				dataset_size=7    ,				batch_size=2    ,				process_0_expected_batch_sizes=[2, 2]    ,				process_1_expected_batch_sizes=[2, 2]    ,				)



def 					_A			() ->       List[Any]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  str    						=		create_accelerator(even_batches=__SCREAMING_SNAKE_CASE							)

    verify_dataloader_batch_sizes(
        __SCREAMING_SNAKE_CASE    ,				dataset_size=3    ,				batch_size=1    ,				process_0_expected_batch_sizes=[1, 1]    ,				process_1_expected_batch_sizes=[1]    ,				)

    verify_dataloader_batch_sizes(
        __SCREAMING_SNAKE_CASE    ,				dataset_size=7    ,				batch_size=2    ,				process_0_expected_batch_sizes=[2, 2]    ,				process_1_expected_batch_sizes=[2, 1]    ,				)



def 					_A			() ->       Optional[int]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  Tuple    						=		create_accelerator(even_batches=__SCREAMING_SNAKE_CASE							)

    SCREAMING_SNAKE_CASE_      :  Any    						=		torch.nn.Linear(1    ,				1							)
    SCREAMING_SNAKE_CASE_      :  Dict    						=		accelerator.prepare(__SCREAMING_SNAKE_CASE							)

    SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		create_dataloader(__SCREAMING_SNAKE_CASE    ,				dataset_size=3    ,				batch_size=1							)

    SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		[]
    with accelerator.join_uneven_inputs([ddp_model]							):
        for batch_idx, batch in enumerate(__SCREAMING_SNAKE_CASE							):
            SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		ddp_model(batch[0].float()							)
            SCREAMING_SNAKE_CASE_      :  int    						=		output.sum()
            loss.backward()
            batch_idxs.append(__SCREAMING_SNAKE_CASE							)

    accelerator.wait_for_everyone()

    if accelerator.process_index == 0:
        assert batch_idxs == [0, 1]
    elif accelerator.process_index == 1:
        assert batch_idxs == [0]



def 					_A			(__a							) ->       List[str]:
    """simple docstring"""






    with warnings.catch_warnings(record=__SCREAMING_SNAKE_CASE							) as w:
        with accelerator.join_uneven_inputs([Mock()]							):
            pass

        assert issubclass(w[-1].category    ,				__SCREAMING_SNAKE_CASE							)
        assert "only supported for multi-GPU" in str(w[-1].message							)



def 					_A			() ->       List[Any]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		True
    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		False
    SCREAMING_SNAKE_CASE_      :  List[str]    						=		create_accelerator(even_batches=__SCREAMING_SNAKE_CASE							)
    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		torch.nn.Linear(1    ,				1							)
    SCREAMING_SNAKE_CASE_      :  List[str]    						=		accelerator.prepare(__SCREAMING_SNAKE_CASE							)
    SCREAMING_SNAKE_CASE_      :  Tuple    						=		create_dataloader(__SCREAMING_SNAKE_CASE    ,				dataset_size=3    ,				batch_size=1							)
    SCREAMING_SNAKE_CASE_      :  int    						=		create_dataloader(__SCREAMING_SNAKE_CASE    ,				dataset_size=3    ,				batch_size=1							)

    with accelerator.join_uneven_inputs([ddp_model]    ,				even_batches=__SCREAMING_SNAKE_CASE							):
        SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		train_dl.batch_sampler.even_batches
        SCREAMING_SNAKE_CASE_      :  str    						=		valid_dl.batch_sampler.even_batches

    assert train_dl_overridden_value == overridden_even_batches
    assert valid_dl_overridden_value == overridden_even_batches
    assert train_dl.batch_sampler.even_batches == default_even_batches
    assert valid_dl.batch_sampler.even_batches == default_even_batches



def 					_A			() ->       List[Any]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  List[Any]    						=		True
    SCREAMING_SNAKE_CASE_      :  Tuple    						=		False
    SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		create_accelerator(even_batches=__SCREAMING_SNAKE_CASE							)
    SCREAMING_SNAKE_CASE_      :  List[Any]    						=		torch.nn.Linear(1    ,				1							)
    SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		accelerator.prepare(__SCREAMING_SNAKE_CASE							)
    create_dataloader(__SCREAMING_SNAKE_CASE    ,				dataset_size=3    ,				batch_size=1    ,				iterable=__SCREAMING_SNAKE_CASE							)
    SCREAMING_SNAKE_CASE_      :  str    						=		create_dataloader(__SCREAMING_SNAKE_CASE    ,				dataset_size=3    ,				batch_size=1							)

    with warnings.catch_warnings():
        warnings.filterwarnings('''ignore'''							)
        try:
            with accelerator.join_uneven_inputs([ddp_model]    ,				even_batches=__SCREAMING_SNAKE_CASE							):
                SCREAMING_SNAKE_CASE_      :  Dict    						=		batch_dl.batch_sampler.even_batches
        except AttributeError:
            # ensure attribute error is not raised when processing iterable dl
            raise AssertionError

    assert batch_dl_overridden_value == overridden_even_batches
    assert batch_dl.batch_sampler.even_batches == default_even_batches



def 					_A			() ->       Optional[int]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  Any    						=		create_accelerator()
    SCREAMING_SNAKE_CASE_      :  str    						=		torch.nn.Linear(1    ,				1							)
    SCREAMING_SNAKE_CASE_      :  Tuple    						=		accelerator.prepare(__SCREAMING_SNAKE_CASE							)
    create_dataloader(__SCREAMING_SNAKE_CASE    ,				dataset_size=3    ,				batch_size=1    ,				iterable=__SCREAMING_SNAKE_CASE							)

    with warnings.catch_warnings(record=__SCREAMING_SNAKE_CASE							) as w:
        with accelerator.join_uneven_inputs([ddp_model]    ,				even_batches=__SCREAMING_SNAKE_CASE							):
            pass

        assert issubclass(w[-1].category    ,				__SCREAMING_SNAKE_CASE							)
        assert "only supported for map-style datasets" in str(w[-1].message							)



def 					_A			() ->       List[Any]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  str    						=		create_accelerator()

    accelerator.print('''Test that even_batches variable ensures uniform batches across processes'''							)
    test_default_ensures_even_batch_sizes()

    accelerator.print('''Run tests with even_batches disabled'''							)
    test_can_disable_even_batches()

    accelerator.print('''Test joining uneven inputs'''							)
    test_can_join_uneven_inputs()

    accelerator.print('''Test overriding even_batches when joining uneven inputs'''							)
    test_join_can_override_even_batches()

    accelerator.print('''Test overriding even_batches for mixed dataloader types'''							)
    test_join_can_override_for_mixed_type_dataloaders()

    accelerator.print('''Test overriding even_batches raises a warning for iterable dataloaders'''							)
    test_join_raises_warning_for_iterable_when_overriding_even_batches()

    accelerator.print('''Test join with non DDP distributed raises warning'''							)
    SCREAMING_SNAKE_CASE_      :  Tuple    						=		accelerator.state.distributed_type
    SCREAMING_SNAKE_CASE_      :  Dict    						=		DistributedType.FSDP
    test_join_raises_warning_for_non_ddp_distributed(__SCREAMING_SNAKE_CASE							)
    SCREAMING_SNAKE_CASE_      :  List[Any]    						=		original_state


if __name__ == "__main__":
  main()





"""simple docstring"""


import argparse

import torch

from transformers import GPTaConfig, GPTaModel, load_tf_weights_in_gpta
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()






def     a__    (    __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )	->	Optional[int]:
		# Construct model
		if gpta_config_file == "":
				__lowerCAmelCase:       Optional[int]          =			GPTaConfig()
		else:
				__lowerCAmelCase:       List[str]          =			GPTaConfig.from_json_file(__SCREAMING_SNAKE_CASE )
		__lowerCAmelCase:       Optional[int]          =			GPTaModel(__SCREAMING_SNAKE_CASE )

		# Load weights from numpy
		load_tf_weights_in_gpta(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )

		# Save pytorch-model
		__lowerCAmelCase:       str          =			pytorch_dump_folder_path + "/" + WEIGHTS_NAME
		__lowerCAmelCase:       List[Any]          =			pytorch_dump_folder_path + "/" + CONFIG_NAME
		print(F"Save PyTorch model to {pytorch_weights_dump_path}" )
		torch.save(model.state_dict()					, __SCREAMING_SNAKE_CASE )
		print(F"Save configuration file to {pytorch_config_dump_path}" )
		with open(__SCREAMING_SNAKE_CASE					, "w"					, encoding="utf-8" ) as f:
				f.write(config.to_json_string() )


if __name__ == "__main__":
			__A						   =			argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    "--gpt2_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
			)
			parser.add_argument(
			    "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
			)
			parser.add_argument(
			    "--gpt2_config_file",
			    default="",
			    type=str,
			    help=(
			        "An optional config json file corresponding to the pre-trained OpenAI model. \n"
			        "This specifies the model architecture."
			    ),
			)
			__A						   =			parser.parse_args()
			convert_gpta_checkpoint_to_pytorch(args.gpta_checkpoint_path, args.gpta_config_file, args.pytorch_dump_folder_path)


