




import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


UpperCamelCase								=     logging.getLogger(__name__)

require_version('pytorch_lightning>=1.0.4')

UpperCamelCase								=     {
    'base': AutoModel,
    'sequence-classification': AutoModelForSequenceClassification,
    'question-answering': AutoModelForQuestionAnswering,
    'pretraining': AutoModelForPreTraining,
    'token-classification': AutoModelForTokenClassification,
    'language-modeling': AutoModelWithLMHead,
    'summarization': AutoModelForSeqaSeqLM,
    'translation': AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
UpperCamelCase								=     {
    'linear': get_linear_schedule_with_warmup,
    'cosine': get_cosine_schedule_with_warmup,
    'cosine_w_restarts': get_cosine_with_hard_restarts_schedule_with_warmup,
    'polynomial': get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
UpperCamelCase								=     sorted(arg_to_scheduler.keys())
UpperCamelCase								=     '{' + ', '.join(arg_to_scheduler_choices) + '}'





class  __lowerCamelCase       (    pl.LightningModule    ):




    """simple docstring"""

    def __init__(       self     :     int   ,	SCREAMING_SNAKE_CASE__     :     argparse.Namespace   ,	SCREAMING_SNAKE_CASE__     :     Optional[int]=None   ,	SCREAMING_SNAKE_CASE__     :     Union[str, Any]="base"   ,	SCREAMING_SNAKE_CASE__     :     str=None   ,	SCREAMING_SNAKE_CASE__     :     Dict=None   ,	SCREAMING_SNAKE_CASE__     :     Optional[Any]=None   ,	**SCREAMING_SNAKE_CASE__     :     Optional[int]   ,	)      ->       List[str]:
         super().__init__()
         # TODO: move to self.save_hyperparameters()
         # self.save_hyperparameters()
         # can also expand arguments into trainer signature for easier reading

         self.save_hyperparameters(SCREAMING_SNAKE_CASE__   )
         lowerCAmelCase__		= 0
         lowerCAmelCase__		= Path(self.hparams.output_dir   )
         lowerCAmelCase__		= self.hparams.cache_dir if self.hparams.cache_dir else None
         if config is None:
              lowerCAmelCase__		= AutoConfig.from_pretrained(
                  self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path   ,	**({"num_labels": num_labels} if num_labels is not None else {})   ,	cache_dir=SCREAMING_SNAKE_CASE__   ,	**SCREAMING_SNAKE_CASE__   ,	)
         else:
              lowerCAmelCase__		= config

         lowerCAmelCase__		= ("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
         for p in extra_model_params:
              if getattr(self.hparams   ,	SCREAMING_SNAKE_CASE__   ,	SCREAMING_SNAKE_CASE__   ):
                   assert hasattr(self.config   ,	SCREAMING_SNAKE_CASE__   ), f'model config doesn\'t have a `{p}` attribute'
                   setattr(self.config   ,	SCREAMING_SNAKE_CASE__   ,	getattr(self.hparams   ,	SCREAMING_SNAKE_CASE__   )   )

         if tokenizer is None:
              lowerCAmelCase__		= AutoTokenizer.from_pretrained(
                  self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path   ,	cache_dir=SCREAMING_SNAKE_CASE__   ,	)
         else:
              lowerCAmelCase__		= tokenizer
         lowerCAmelCase__		= MODEL_MODES[mode]
         if model is None:
              lowerCAmelCase__		= self.model_type.from_pretrained(
                  self.hparams.model_name_or_path   ,	from_tf=bool(".ckpt" in self.hparams.model_name_or_path   )   ,	config=self.config   ,	cache_dir=SCREAMING_SNAKE_CASE__   ,	)
         else:
              lowerCAmelCase__		= model


    def    a      (       self     :     Tuple   ,	*SCREAMING_SNAKE_CASE__     :     List[str]   ,	**SCREAMING_SNAKE_CASE__     :     List[str]   )      ->       List[Any]:
         lowerCAmelCase__		= self.model_type.from_pretrained(*SCREAMING_SNAKE_CASE__   ,	**SCREAMING_SNAKE_CASE__   )


    def    a      (       self     :     Union[str, Any]   )      ->       Optional[Any]:
         lowerCAmelCase__		= arg_to_scheduler[self.hparams.lr_scheduler]
         lowerCAmelCase__		= get_schedule_func(
             self.opt   ,	num_warmup_steps=self.hparams.warmup_steps   ,	num_training_steps=self.total_steps()   )
         lowerCAmelCase__		= {"scheduler": scheduler, "interval": "step", "frequency": 1}
         return scheduler


    def    a      (       self     :     Tuple   )      ->       str:
         lowerCAmelCase__		= self.model
         lowerCAmelCase__		= ["bias", "LayerNorm.weight"]
         lowerCAmelCase__		= [
             {
                 "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay   )],
                 "weight_decay": self.hparams.weight_decay,
             },
             {
                 "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay   )],
                 "weight_decay": 0.0,
             },
         ]
         if self.hparams.adafactor:
              lowerCAmelCase__		= Adafactor(
                  SCREAMING_SNAKE_CASE__   ,	lr=self.hparams.learning_rate   ,	scale_parameter=SCREAMING_SNAKE_CASE__   ,	relative_step=SCREAMING_SNAKE_CASE__   )

         else:
              lowerCAmelCase__		= AdamW(
                  SCREAMING_SNAKE_CASE__   ,	lr=self.hparams.learning_rate   ,	eps=self.hparams.adam_epsilon   )
         lowerCAmelCase__		= optimizer

         lowerCAmelCase__		= self.get_lr_scheduler()

         return [optimizer], [scheduler]


    def    a      (       self     :     Optional[Any]   ,	SCREAMING_SNAKE_CASE__     :     Union[str, Any]   ,	SCREAMING_SNAKE_CASE__     :     Optional[int]   )      ->       Union[str, Any]:
         return self.validation_step(SCREAMING_SNAKE_CASE__   ,	SCREAMING_SNAKE_CASE__   )


    def    a      (       self     :     Tuple   ,	SCREAMING_SNAKE_CASE__     :     Dict   )      ->       str:
         return self.validation_end(SCREAMING_SNAKE_CASE__   )


    def    a      (       self     :     Any   )      ->       int:
         lowerCAmelCase__		= max(1   ,	self.hparams.gpus   )  # TODO: consider num_tpu_cores
         lowerCAmelCase__		= self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
         return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs


    def    a      (       self     :     List[str]   ,	SCREAMING_SNAKE_CASE__     :     Any   )      ->       List[str]:
         if mode == "test":
              lowerCAmelCase__		= len(self.test_dataloader().dataset   )
         else:
              lowerCAmelCase__		= self.get_dataloader("train"   ,	self.hparams.train_batch_size   ,	shuffle=SCREAMING_SNAKE_CASE__   )
              lowerCAmelCase__		= len(self.train_dataloader().dataset   )


    def    a      (       self     :     Dict   ,	SCREAMING_SNAKE_CASE__     :     str   ,	SCREAMING_SNAKE_CASE__     :     int   ,	SCREAMING_SNAKE_CASE__     :     bool = False   )      ->       List[str]:
         raise NotImplementedError("You must implement this for your task"   )


    def    a      (       self     :     List[str]   )      ->       Optional[Any]:
         return self.train_loader


    def    a      (       self     :     Dict   )      ->       Dict:
         return self.get_dataloader("dev"   ,	self.hparams.eval_batch_size   ,	shuffle=SCREAMING_SNAKE_CASE__   )


    def    a      (       self     :     Optional[Any]   )      ->       List[str]:
         return self.get_dataloader("test"   ,	self.hparams.eval_batch_size   ,	shuffle=SCREAMING_SNAKE_CASE__   )


    def    a      (       self     :     Union[str, Any]   ,	SCREAMING_SNAKE_CASE__     :     Any   )      ->       Dict:
         return os.path.join(
             self.hparams.data_dir   ,	"cached_{}_{}_{}".format(
                 SCREAMING_SNAKE_CASE__   ,	list(filter(SCREAMING_SNAKE_CASE__   ,	self.hparams.model_name_or_path.split("/"   )   )   ).pop()   ,	str(self.hparams.max_seq_length   )   ,	)   ,	)


    @pl.utilities.rank_zero_only
    def    a      (       self     :     Union[str, Any]   ,	SCREAMING_SNAKE_CASE__     :     Dict[str, Any]   )      ->       None:
         lowerCAmelCase__		= self.output_dir.joinpath("best_tfmr"   )
         lowerCAmelCase__		= self.step_count
         self.model.save_pretrained(SCREAMING_SNAKE_CASE__   )
         self.tokenizer.save_pretrained(SCREAMING_SNAKE_CASE__   )


    @staticmethod
    def    a      (       SCREAMING_SNAKE_CASE__     :     Optional[Any]   ,	SCREAMING_SNAKE_CASE__     :     Any   )      ->       Optional[Any]:
         parser.add_argument(
             "--model_name_or_path"   ,	default=SCREAMING_SNAKE_CASE__   ,	type=SCREAMING_SNAKE_CASE__   ,	required=SCREAMING_SNAKE_CASE__   ,	help="Path to pretrained model or model identifier from huggingface.co/models"   ,	)
         parser.add_argument(
             "--config_name"   ,	default=""   ,	type=SCREAMING_SNAKE_CASE__   ,	help="Pretrained config name or path if not the same as model_name"   )
         parser.add_argument(
             "--tokenizer_name"   ,	default=SCREAMING_SNAKE_CASE__   ,	type=SCREAMING_SNAKE_CASE__   ,	help="Pretrained tokenizer name or path if not the same as model_name"   ,	)
         parser.add_argument(
             "--cache_dir"   ,	default=""   ,	type=SCREAMING_SNAKE_CASE__   ,	help="Where do you want to store the pre-trained models downloaded from huggingface.co"   ,	)
         parser.add_argument(
             "--encoder_layerdrop"   ,	type=SCREAMING_SNAKE_CASE__   ,	help="Encoder layer dropout probability (Optional). Goes into model.config"   ,	)
         parser.add_argument(
             "--decoder_layerdrop"   ,	type=SCREAMING_SNAKE_CASE__   ,	help="Decoder layer dropout probability (Optional). Goes into model.config"   ,	)
         parser.add_argument(
             "--dropout"   ,	type=SCREAMING_SNAKE_CASE__   ,	help="Dropout probability (Optional). Goes into model.config"   ,	)
         parser.add_argument(
             "--attention_dropout"   ,	type=SCREAMING_SNAKE_CASE__   ,	help="Attention dropout probability (Optional). Goes into model.config"   ,	)
         parser.add_argument("--learning_rate"   ,	default=5e-5   ,	type=SCREAMING_SNAKE_CASE__   ,	help="The initial learning rate for Adam."   )
         parser.add_argument(
             "--lr_scheduler"   ,	default="linear"   ,	choices=SCREAMING_SNAKE_CASE__   ,	metavar=SCREAMING_SNAKE_CASE__   ,	type=SCREAMING_SNAKE_CASE__   ,	help="Learning rate scheduler"   ,	)
         parser.add_argument("--weight_decay"   ,	default=0.0   ,	type=SCREAMING_SNAKE_CASE__   ,	help="Weight decay if we apply some."   )
         parser.add_argument("--adam_epsilon"   ,	default=1e-8   ,	type=SCREAMING_SNAKE_CASE__   ,	help="Epsilon for Adam optimizer."   )
         parser.add_argument("--warmup_steps"   ,	default=0   ,	type=SCREAMING_SNAKE_CASE__   ,	help="Linear warmup over warmup_steps."   )
         parser.add_argument("--num_workers"   ,	default=4   ,	type=SCREAMING_SNAKE_CASE__   ,	help="kwarg passed to DataLoader"   )
         parser.add_argument("--num_train_epochs"   ,	dest="max_epochs"   ,	default=3   ,	type=SCREAMING_SNAKE_CASE__   )
         parser.add_argument("--train_batch_size"   ,	default=32   ,	type=SCREAMING_SNAKE_CASE__   )
         parser.add_argument("--eval_batch_size"   ,	default=32   ,	type=SCREAMING_SNAKE_CASE__   )
         parser.add_argument("--adafactor"   ,	action="store_true"   )






class  __lowerCamelCase       (    pl.Callback    ):




    """simple docstring"""

    def    a      (       self     :     Optional[Any]   ,	SCREAMING_SNAKE_CASE__     :     Any   ,	SCREAMING_SNAKE_CASE__     :     Optional[int]   )      ->       List[Any]:
         lowerCAmelCase__		= trainer.lr_schedulers[0]["scheduler"]
         lowerCAmelCase__		= {f'lr_group_{i}': lr for i, lr in enumerate(lr_scheduler.get_lr()   )}
         pl_module.logger.log_metrics(SCREAMING_SNAKE_CASE__   )


    def    a      (       self     :     List[Any]   ,	SCREAMING_SNAKE_CASE__     :     pl.Trainer   ,	SCREAMING_SNAKE_CASE__     :     pl.LightningModule   )      ->       Any:
         rank_zero_info("***** Validation results *****"   )
         lowerCAmelCase__		= trainer.callback_metrics
         # Log results
         for key in sorted(SCREAMING_SNAKE_CASE__   ):
              if key not in ["log", "progress_bar"]:
                   rank_zero_info("{} = {}\n".format(SCREAMING_SNAKE_CASE__   ,	str(metrics[key]   )   )   )


    def    a      (       self     :     Tuple   ,	SCREAMING_SNAKE_CASE__     :     pl.Trainer   ,	SCREAMING_SNAKE_CASE__     :     pl.LightningModule   )      ->       int:
         rank_zero_info("***** Test results *****"   )
         lowerCAmelCase__		= trainer.callback_metrics
         # Log and save results to file
         lowerCAmelCase__		= os.path.join(pl_module.hparams.output_dir   ,	"test_results.txt"   )
         with open(SCREAMING_SNAKE_CASE__   ,	"w"   ) as writer:
              for key in sorted(SCREAMING_SNAKE_CASE__   ):
                   if key not in ["log", "progress_bar"]:
                        rank_zero_info("{} = {}\n".format(SCREAMING_SNAKE_CASE__   ,	str(metrics[key]   )   )   )
                        writer.write("{} = {}\n".format(SCREAMING_SNAKE_CASE__   ,	str(metrics[key]   )   )   )






def 							_A   (						lowerCAmelCase_     :					int   ,				lowerCAmelCase_     :					str			):
     """simple docstring"""

     parser.add_argument(
         "--output_dir"   ,				default=lowerCAmelCase_   ,				type=lowerCAmelCase_   ,				required=lowerCAmelCase_   ,				help="The output directory where the model predictions and checkpoints will be written."   ,				)
     parser.add_argument(
         "--fp16"   ,				action="store_true"   ,				help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"   ,				)

     parser.add_argument(
         "--fp16_opt_level"   ,				type=lowerCAmelCase_   ,				default="O2"   ,				help=(
             "For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
             "See details at https://nvidia.github.io/apex/amp.html"
         )   ,				)
     parser.add_argument("--n_tpu_cores"   ,				dest="tpu_cores"   ,				type=lowerCAmelCase_			)
     parser.add_argument("--max_grad_norm"   ,				dest="gradient_clip_val"   ,				default=1.0   ,				type=lowerCAmelCase_   ,				help="Max gradient norm"			)
     parser.add_argument("--do_train"   ,				action="store_true"   ,				help="Whether to run training."			)
     parser.add_argument("--do_predict"   ,				action="store_true"   ,				help="Whether to run predictions on the test set."			)
     parser.add_argument(
         "--gradient_accumulation_steps"   ,				dest="accumulate_grad_batches"   ,				type=lowerCAmelCase_   ,				default=1   ,				help="Number of updates steps to accumulate before performing a backward/update pass."   ,				)
     parser.add_argument("--seed"   ,				type=lowerCAmelCase_   ,				default=42   ,				help="random seed for initialization"			)
     parser.add_argument(
         "--data_dir"   ,				default=lowerCAmelCase_   ,				type=lowerCAmelCase_   ,				required=lowerCAmelCase_   ,				help="The input data dir. Should contain the training files for the CoNLL-2003 NER task."   ,				)






def 							_A   (						lowerCAmelCase_     :					BaseTransformer   ,				lowerCAmelCase_     :					argparse.Namespace   ,				lowerCAmelCase_     :					Dict=None   ,				lowerCAmelCase_     :					str=True   ,				lowerCAmelCase_     :					Tuple=[]   ,				lowerCAmelCase_     :					Any=None   ,				lowerCAmelCase_     :					Any=None   ,				**lowerCAmelCase_     :					List[Any]   ,				):
     """simple docstring"""

     pl.seed_everything(args.seed			)

     # init model
     lowerCAmelCase__		= Path(model.hparams.output_dir			)
     odir.mkdir(exist_ok=lowerCAmelCase_			)

     # add custom checkpoints
     if checkpoint_callback is None:
          lowerCAmelCase__		= pl.callbacks.ModelCheckpoint(
              filepath=args.output_dir   ,				prefix="checkpoint"   ,				monitor="val_loss"   ,				mode="min"   ,				save_top_k=1			)
     if early_stopping_callback:
          extra_callbacks.append(lowerCAmelCase_			)
     if logging_callback is None:
          lowerCAmelCase__		= LoggingCallback()

     lowerCAmelCase__		= {}

     # TODO: remove with PyTorch 1.6 since pl uses native amp
     if args.fpaa:
          lowerCAmelCase__		= 16
          lowerCAmelCase__		= args.fpaa_opt_level

     if args.gpus > 1:
          lowerCAmelCase__		= "ddp"

     lowerCAmelCase__		= args.accumulate_grad_batches
     lowerCAmelCase__		= extra_train_kwargs.get("accelerator"   ,				lowerCAmelCase_			)
     lowerCAmelCase__		= extra_train_kwargs.get("profiler"   ,				lowerCAmelCase_			)

     lowerCAmelCase__		= pl.Trainer.from_argparse_args(
         lowerCAmelCase_   ,				weights_summary=lowerCAmelCase_   ,				callbacks=[logging_callback] + extra_callbacks   ,				logger=lowerCAmelCase_   ,				checkpoint_callback=lowerCAmelCase_   ,				**lowerCAmelCase_   ,				)

     if args.do_train:
          trainer.fit(lowerCAmelCase_			)

     return trainer






import argparse

import torch
from flax.training.checkpoints import restore_checkpoint

from transformers import FNetConfig, FNetForPreTraining
from transformers.utils import logging


logging.set_verbosity_info()



def 							_A   (						lowerCAmelCase_     :					Dict   ,				lowerCAmelCase_     :					List[str]   ,				lowerCAmelCase_     :					int			):
     """simple docstring"""

     lowerCAmelCase__		= FNetConfig.from_json_file(lowerCAmelCase_			)
     print(F'Building PyTorch model from configuration: {config}'			)
     lowerCAmelCase__		= FNetForPreTraining(lowerCAmelCase_			)

     lowerCAmelCase__		= restore_checkpoint(lowerCAmelCase_   ,				lowerCAmelCase_			)
     lowerCAmelCase__		= checkpoint_dict["target"]

     # Embeddings
     # Position IDs
     lowerCAmelCase__		= fnet_pretraining_model.state_dict()

     lowerCAmelCase__		= state_dict["fnet.embeddings.position_ids"]
     lowerCAmelCase__		= {"fnet.embeddings.position_ids": position_ids}
     # Embedding Layers
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["encoder"]["embedder"]["word"]["embedding"]			)
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["encoder"]["embedder"]["position"]["embedding"][0]			)
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["encoder"]["embedder"]["type"]["embedding"]			)
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["encoder"]["embedder"]["hidden_mapping_in"]["kernel"]			).T
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["encoder"]["embedder"]["hidden_mapping_in"]["bias"]			)
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["encoder"]["embedder"]["layer_norm"]["scale"]			)
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["encoder"]["embedder"]["layer_norm"]["bias"]			)

     # Encoder Layers
     for layer in range(config.num_hidden_layers			):
          lowerCAmelCase__		= torch.tensor(
              pretrained_model_params["encoder"][F'encoder_{layer}']["mixing_layer_norm"]["scale"]			)
          lowerCAmelCase__		= torch.tensor(
              pretrained_model_params["encoder"][F'encoder_{layer}']["mixing_layer_norm"]["bias"]			)

          lowerCAmelCase__		= torch.tensor(
              pretrained_model_params["encoder"][F'feed_forward_{layer}']["intermediate"]["kernel"]			).T
          lowerCAmelCase__		= torch.tensor(
              pretrained_model_params["encoder"][F'feed_forward_{layer}']["intermediate"]["bias"]			)

          lowerCAmelCase__		= torch.tensor(
              pretrained_model_params["encoder"][F'feed_forward_{layer}']["output"]["kernel"]			).T
          lowerCAmelCase__		= torch.tensor(
              pretrained_model_params["encoder"][F'feed_forward_{layer}']["output"]["bias"]			)

          lowerCAmelCase__		= torch.tensor(
              pretrained_model_params["encoder"][F'encoder_{layer}']["output_layer_norm"]["scale"]			)
          lowerCAmelCase__		= torch.tensor(
              pretrained_model_params["encoder"][F'encoder_{layer}']["output_layer_norm"]["bias"]			)

     # Pooler Layers
     lowerCAmelCase__		= torch.tensor(pretrained_model_params["encoder"]["pooler"]["kernel"]			).T
     lowerCAmelCase__		= torch.tensor(pretrained_model_params["encoder"]["pooler"]["bias"]			)

     # Masked LM Layers
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["predictions_dense"]["kernel"]			).T
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["predictions_dense"]["bias"]			)
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["predictions_layer_norm"]["scale"]			)
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["predictions_layer_norm"]["bias"]			)
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["encoder"]["embedder"]["word"]["embedding"]			)
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["predictions_output"]["output_bias"]			)
     lowerCAmelCase__		= torch.tensor(pretrained_model_params["predictions_output"]["output_bias"]			)

     # Seq Relationship Layers
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["classification"]["output_kernel"]			)
     lowerCAmelCase__		= torch.tensor(
         pretrained_model_params["classification"]["output_bias"]			)

     # Load State Dict
     fnet_pretraining_model.load_state_dict(lowerCAmelCase_			)

     # Save PreTrained
     print(F'Saving pretrained model to {save_path}'			)
     fnet_pretraining_model.save_pretrained(lowerCAmelCase_			)


if __name__ == "__main__":
     UpperCamelCase								=     argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '--flax_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
     )
     parser.add_argument(
         '--fnet_config_file',
         default=None,
         type=str,
         required=True,
         help=(
             'The config json file corresponding to the pre-trained FNet model. \n'
             'This specifies the model architecture.'
         ),
     )
     parser.add_argument('--save_path', default=None, type=str, required=True, help='Path to the output model.')
     UpperCamelCase								=     parser.parse_args()
     convert_flax_checkpoint_to_pytorch(args.flax_checkpoint_path, args.fnet_config_file, args.save_path)
