


"""simple docstring"""
import argparse

import torch

from transformers import ElectraConfig, ElectraForMaskedLM, ElectraForPreTraining, load_tf_weights_in_electra
from transformers.utils import logging


logging.set_verbosity_info()






def 						lowercase				(	_snake_case		: List[Any]	,  _snake_case		: Tuple	,  _snake_case		: Union[str, Any]	,  _snake_case		: Tuple	)	->Optional[Any]:
	"""simple docstring"""





	__snake_case      :		Union[str, Any]         =							ElectraConfig.from_json_file(_snake_case	)
	print(f"""Building PyTorch model from configuration: {config}"""	)

	if discriminator_or_generator == "discriminator":
		__snake_case      :		Optional[Any]         =							ElectraForPreTraining(_snake_case	)
	elif discriminator_or_generator == "generator":
		__snake_case      :		Optional[int]         =							ElectraForMaskedLM(_snake_case	)
	else:
		raise ValueError('''The discriminator_or_generator argument should be either \'discriminator\' or \'generator\''''	)

	# Load weights from tf checkpoint
	load_tf_weights_in_electra(
	    _snake_case	,  _snake_case	,  _snake_case	,  discriminator_or_generator=_snake_case	)

	# Save pytorch-model
	print(f"""Save PyTorch model to {pytorch_dump_path}"""	)
	torch.save(model.state_dict()	,  _snake_case	)


if __name__ == "__main__":
			SCREAMING_SNAKE_CASE					:       List[str]							   =							argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    """--tf_checkpoint_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
			)
			parser.add_argument(
			    """--config_file""",
			    default=None,
			    type=str,
			    required=True,
			    help="""The config json file corresponding to the pre-trained model. \nThis specifies the model architecture.""",
			)
			parser.add_argument(
			    """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
			)
			parser.add_argument(
			    """--discriminator_or_generator""",
			    default=None,
			    type=str,
			    required=True,
			    help=(
			        """Whether to export the generator or the discriminator. Should be a string, either 'discriminator' or """
			        """'generator'."""
			    ),
			)
			SCREAMING_SNAKE_CASE					:       str							   =							parser.parse_args()
			convert_tf_checkpoint_to_pytorch(
			    args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path, args.discriminator_or_generator
			)












'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


snake_case_  :							Any            =	{
    "configuration_flava": [
        "FLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "FlavaConfig",
        "FlavaImageCodebookConfig",
        "FlavaImageConfig",
        "FlavaMultimodalConfig",
        "FlavaTextConfig",
    ],
}

try:
					if not is_vision_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_  :							List[str]            =	["FlavaFeatureExtractor"]
					snake_case_  :							int            =	["FlavaImageProcessor"]
					snake_case_  :							Tuple            =	["FlavaProcessor"]

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_  :							Optional[int]            =	[
					    "FLAVA_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "FlavaForPreTraining",
					    "FlavaImageCodebook",
					    "FlavaImageModel",
					    "FlavaModel",
					    "FlavaMultimodalModel",
					    "FlavaPreTrainedModel",
					    "FlavaTextModel",
					]

if TYPE_CHECKING:
					from .configuration_flava import (
					    FLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP,
					    FlavaConfig,
					    FlavaImageCodebookConfig,
					    FlavaImageConfig,
					    FlavaMultimodalConfig,
					    FlavaTextConfig,
					)

					try:
										if not is_vision_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .feature_extraction_flava import FlavaFeatureExtractor
										from .image_processing_flava import FlavaImageProcessor
										from .processing_flava import FlavaProcessor

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_flava import (
										    FLAVA_PRETRAINED_MODEL_ARCHIVE_LIST,
										    FlavaForPreTraining,
										    FlavaImageCodebook,
										    FlavaImageModel,
										    FlavaModel,
										    FlavaMultimodalModel,
										    FlavaPreTrainedModel,
										    FlavaTextModel,
										)

else:
					import sys

					snake_case_  :							Any            =	_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





