'''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 __lowerCAmelCase				(							UpperCamelCase__       ,						UpperCamelCase__       ,						UpperCamelCase__       ,						UpperCamelCase__       )			->					Any:
    # Initialise PyTorch model
    __lowerCamelCase						=						ElectraConfig.from_json_file(UpperCamelCase__       )
    print(f"""Building PyTorch model from configuration: {config}"""       )

    if discriminator_or_generator == "discriminator":
        __lowerCamelCase						=						ElectraForPreTraining(UpperCamelCase__       )
    elif discriminator_or_generator == "generator":
        __lowerCamelCase						=						ElectraForMaskedLM(UpperCamelCase__       )
    else:
        raise ValueError('''The discriminator_or_generator argument should be either \'discriminator\' or \'generator\''''       )

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

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


if __name__ == "__main__":
  __UpperCAmelCase  							=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'."
      ),
  )
  __UpperCAmelCase  							=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,
)


__UpperCAmelCase  							={
    "configuration_rwkv": ["RWKV_PRETRAINED_CONFIG_ARCHIVE_MAP", "RwkvConfig", "RwkvOnnxConfig"],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase  							=[
      "RWKV_PRETRAINED_MODEL_ARCHIVE_LIST",
      "RwkvForCausalLM",
      "RwkvModel",
      "RwkvPreTrainedModel",
  ]


if TYPE_CHECKING:
  from .configuration_rwkv import RWKV_PRETRAINED_CONFIG_ARCHIVE_MAP, RwkvConfig, RwkvOnnxConfig

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_rwkv import (
        RWKV_PRETRAINED_MODEL_ARCHIVE_LIST,
        RwkvForCausalLM,
        RwkvModel,
        RwkvPreTrainedModel,
    )
else:
  import sys

  __UpperCAmelCase  							=_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
