
"""simple docstring"""




import argparse
import gc
import os
import threading

import evaluate
import psutil
import torch
from datasets import load_dataset
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig
from torch.utils.data import DataLoader
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    get_linear_schedule_with_warmup,
    set_seed,
)

from accelerate import Accelerator, DistributedType, FullyShardedDataParallelPlugin


########################################################################
# This is a fully working simple example to use Accelerate
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#   - FSDP
#
# This example also demonstrates the checkpointing and sharding capabilities
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


_lowerCamelCase        =1_6
_lowerCamelCase        =3_2





def    _a (						lowerCamelCase							):
 return int(x / 2**20							)





class 							A__		:
      def __enter__(				self				):
       gc.collect()
       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
       lowerCamelCase		:     Union[str, Any]          =				torch.cuda.memory_allocated()
       lowerCamelCase		:     List[str]          =				psutil.Process()

       lowerCamelCase		:     Any          =				self.cpu_mem_used()
       lowerCamelCase		:     Union[str, Any]          =				True
       lowerCamelCase		:     Any          =				threading.Thread(target=self.peak_monitor_func				)
       lowerCamelCase		:     Any          =				True
       peak_monitor_thread.start()
       return self
      def 						UpperCamelCase__							(				self				):
       return self.process.memory_info().rss
      def 						UpperCamelCase__							(				self				):
       lowerCamelCase		:     Any          =				-1

       while True:
        lowerCamelCase		:     Dict          =				max(self.cpu_mem_used()   ,  self.cpu_peak				)

        # can't sleep or will not catch the peak right (this comment is here on purpose)
        # time.sleep(0.001) # 1msec

        if not self.peak_monitoring:
         break






      def __exit__(				self   ,  *__magic_name__				):
       lowerCamelCase		:     Dict          =				False

       gc.collect()
       torch.cuda.empty_cache()
       lowerCamelCase		:     Optional[int]          =				torch.cuda.memory_allocated()
       lowerCamelCase		:     Union[str, Any]          =				torch.cuda.max_memory_allocated()
       lowerCamelCase		:     Optional[int]          =				bamb(self.end - self.begin				)
       lowerCamelCase		:     Optional[Any]          =				bamb(self.peak - self.begin				)

       lowerCamelCase		:     int          =				self.cpu_mem_used()
       lowerCamelCase		:     Dict          =				bamb(self.cpu_end - self.cpu_begin				)
       lowerCamelCase		:     Optional[Any]          =				bamb(self.cpu_peak - self.cpu_begin				)
       # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")


# For testing only
if os.environ.get("""TESTING_MOCKED_DATALOADERS""", None) == "1":
      from accelerate.test_utils.training import mocked_dataloaders

      _lowerCamelCase        =mocked_dataloaders  # noqa: F811





def    _a (						lowerCamelCase,     lowerCamelCase							):
 # For testing only
 if os.environ.get("""TESTING_MOCKED_DATALOADERS""",     _lowerCamelCase							) == "1":
  lowerCamelCase		:     Dict          =				2

 # New Code #
 # Pass the advanced FSDP settings not part of the accelerate config by creating fsdp_plugin
 lowerCamelCase		:     Dict          =				FullyShardedDataParallelPlugin(
     state_dict_config=FullStateDictConfig(offload_to_cpu=_lowerCamelCase,     ranka_only=_lowerCamelCase							),     optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=_lowerCamelCase,     ranka_only=_lowerCamelCase							),     )

 # Initialize accelerator
 if args.with_tracking:
  lowerCamelCase		:     List[str]          =				Accelerator(
      cpu=args.cpu,     mixed_precision=args.mixed_precision,     log_with="""wandb""",     project_dir=args.logging_dir,     fsdp_plugin=_lowerCamelCase,     )
 else:
  lowerCamelCase		:     Optional[Any]          =				Accelerator(fsdp_plugin=_lowerCamelCase							)
 accelerator.print(accelerator.distributed_type							)

 if hasattr(args.checkpointing_steps,     """isdigit"""							):
  if args.checkpointing_steps == "epoch":
   lowerCamelCase		:     Optional[Any]          =				args.checkpointing_steps
  elif args.checkpointing_steps.isdigit():
   lowerCamelCase		:     List[Any]          =				int(args.checkpointing_steps							)
  else:
   raise ValueError(
       F'''Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.'''							)
 else:
  lowerCamelCase		:     Optional[Any]          =				None
 # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
 lowerCamelCase		:     Optional[int]          =				config["""lr"""]
 lowerCamelCase		:     Tuple          =				int(config["""num_epochs"""]							)
 lowerCamelCase		:     int          =				int(config["""seed"""]							)
 lowerCamelCase		:     Tuple          =				int(config["""batch_size"""]							)

 # We need to initialize the trackers we use, and also store our configuration
 if args.with_tracking:
  lowerCamelCase		:     Tuple          =				vars(_lowerCamelCase							)
  accelerator.init_trackers("""fsdp_glue_no_trainer""",     _lowerCamelCase							)

 lowerCamelCase		:     Dict          =				AutoTokenizer.from_pretrained(args.model_name_or_path							)
 lowerCamelCase		:     str          =				load_dataset("""glue""",     """mrpc"""							)
 lowerCamelCase		:     List[Any]          =				evaluate.load("""glue""",     """mrpc"""							)

 def tokenize_function(lowerCamelCase							):
  # max_length=None => use the model max length (it's actually the default)
  lowerCamelCase		:     Any          =				tokenizer(examples["""sentence1"""],     examples["""sentence2"""],     truncation=_lowerCamelCase,     max_length=_lowerCamelCase							)
  return outputs

 # Apply the method we just defined to all the examples in all the splits of the dataset
 # starting with the main process first:
 with accelerator.main_process_first():
  lowerCamelCase		:     List[Any]          =				datasets.map(
      _lowerCamelCase,     batched=_lowerCamelCase,     remove_columns=["""idx""", """sentence1""", """sentence2"""],     )

 # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
 # transformers library
 lowerCamelCase		:     Optional[Any]          =				tokenized_datasets.rename_column("""label""",     """labels"""							)

 # If the batch size is too big we use gradient accumulation
 lowerCamelCase		:     str          =				1
 if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
  lowerCamelCase		:     List[Any]          =				batch_size // MAX_GPU_BATCH_SIZE
  lowerCamelCase		:     Tuple          =				MAX_GPU_BATCH_SIZE

 def collate_fn(lowerCamelCase							):
  # On TPU it's best to pad everything to the same length or training will be very slow.
  lowerCamelCase		:     str          =				128 if accelerator.distributed_type == DistributedType.TPU else None
  # When using mixed precision we want round multiples of 8/16
  if accelerator.mixed_precision == "fp8":
   lowerCamelCase		:     Any          =				16
  elif accelerator.mixed_precision != "no":
   lowerCamelCase		:     Optional[Any]          =				8
  else:
   lowerCamelCase		:     int          =				None

  return tokenizer.pad(
      _lowerCamelCase,     padding="""longest""",     max_length=_lowerCamelCase,     pad_to_multiple_of=_lowerCamelCase,     return_tensors="""pt""",     )

 # Instantiate dataloaders.
 lowerCamelCase		:     List[str]          =				DataLoader(
     tokenized_datasets["""train"""],     shuffle=_lowerCamelCase,     collate_fn=_lowerCamelCase,     batch_size=_lowerCamelCase							)
 lowerCamelCase		:     str          =				DataLoader(
     tokenized_datasets["""validation"""],     shuffle=_lowerCamelCase,     collate_fn=_lowerCamelCase,     batch_size=_lowerCamelCase							)

 set_seed(_lowerCamelCase							)

 # Instantiate the model (we build the model here so that the seed also control new weights initialization)
 lowerCamelCase		:     Optional[int]          =				AutoModelForSequenceClassification.from_pretrained(
     args.model_name_or_path,     return_dict=_lowerCamelCase,     low_cpu_mem_usage=_lowerCamelCase							)

 # New Code #
 # For FSDP feature, it is highly recommended and efficient to prepare the model before creating optimizer
 lowerCamelCase		:     List[str]          =				accelerator.prepare(_lowerCamelCase							)
 accelerator.print(_lowerCamelCase							)

 # Instantiate optimizer
 # New Code #
 # For FSDP feature, at present it doesn't support multiple parameter groups,
 # so we need to create a single parameter group for the whole model
 lowerCamelCase		:     Dict          =				torch.optim.AdamW(params=model.parameters(),     lr=_lowerCamelCase,     weight_decay=2e-4							)

 # Instantiate scheduler
 lowerCamelCase		:     str          =				get_linear_schedule_with_warmup(
     optimizer=_lowerCamelCase,     num_warmup_steps=10,     num_training_steps=(len(_lowerCamelCase							) * num_epochs) // gradient_accumulation_steps,     )

 # New Code #
 # For FSDP feature, prepare everything except the model as we have already prepared the model
 # before creating the optimizer
 # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
 # prepare method.
 lowerCamelCase		:     int          =				accelerator.prepare(
     _lowerCamelCase,     _lowerCamelCase,     _lowerCamelCase,     _lowerCamelCase							)

 lowerCamelCase		:     List[str]          =				0

 # Potentially load in the weights and states from a previous save
 if args.resume_from_checkpoint:
  if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
   accelerator.print(F'''Resumed from checkpoint: {args.resume_from_checkpoint}'''							)
   accelerator.load_state(args.resume_from_checkpoint							)
   lowerCamelCase		:     Tuple          =				os.path.basename(args.resume_from_checkpoint							)
  else:
   # Get the most recent checkpoint
   lowerCamelCase		:     Tuple          =				[f.name for f in os.scandir(os.getcwd()							) if f.is_dir()]
   dirs.sort(key=os.path.getctime							)
   lowerCamelCase		:     Union[str, Any]          =				dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
  # Extract `epoch_{i}` or `step_{i}`
  lowerCamelCase		:     List[Any]          =				os.path.splitext(_lowerCamelCase							)[0]

  if "epoch" in training_difference:
   num_epochs -= int(training_difference.replace("""epoch_""",     """"""							)							)
   lowerCamelCase		:     Union[str, Any]          =				None
  else:
   lowerCamelCase		:     Any          =				int(training_difference.replace("""step_""",     """"""							)							)
   num_epochs -= resume_step // len(_lowerCamelCase							)
   # If resuming by step, we also need to know exactly how far into the DataLoader we went
   lowerCamelCase		:     int          =				(num_epochs * len(_lowerCamelCase							)) - resume_step

  # Now we train the model
 for epoch in range(_lowerCamelCase							):
  # New Code #
  # context manager to track the peak memory usage during the training epoch
  with TorchTracemalloc() as tracemalloc:
   model.train()
   if args.with_tracking:
    lowerCamelCase		:     Tuple          =				0
   for step, batch in enumerate(_lowerCamelCase							):
    # We need to skip steps until we reach the resumed step
    if args.resume_from_checkpoint and epoch == 0:
     if resume_step is not None and step < resume_step:
      pass
          # We could avoid this line since we set the accelerator with `device_placement=True`.
    batch.to(accelerator.device							)
    lowerCamelCase		:     Optional[int]          =				model(**_lowerCamelCase							)
    lowerCamelCase		:     Optional[Any]          =				outputs.loss
    lowerCamelCase		:     Union[str, Any]          =				loss / gradient_accumulation_steps
    # We keep track of the loss at each epoch
    if args.with_tracking:
     total_loss += loss.detach().float()
    accelerator.backward(_lowerCamelCase							)
    if step % gradient_accumulation_steps == 0:
     optimizer.step()
     lr_scheduler.step()
     optimizer.zero_grad()
     # accelerator.print(lr_scheduler.get_lr())

    overall_step += 1

    if isinstance(_lowerCamelCase,     _lowerCamelCase							):
     lowerCamelCase		:     str          =				F'''step_{overall_step}'''
     if overall_step % checkpointing_steps == 0:
      if args.output_dir is not None:
       lowerCamelCase		:     Optional[Any]          =				os.path.join(args.output_dir,     _lowerCamelCase							)
      accelerator.save_state(_lowerCamelCase							)
    # New Code #
    # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
  accelerator.print("""Memory before entering the train : {}""".format(bamb(tracemalloc.begin							)							)							)
  accelerator.print("""Memory consumed at the end of the train (end-begin): {}""".format(tracemalloc.used							)							)
  accelerator.print("""Peak Memory consumed during the train (max-begin): {}""".format(tracemalloc.peaked							)							)
  accelerator.print(
      """Total Peak Memory consumed during the train (max): {}""".format(
          tracemalloc.peaked + bamb(tracemalloc.begin							)							)							)
  # Logging the peak memory usage of the GPU to the tracker
  if args.with_tracking:
   accelerator.log(
       {
           """train_total_peak_memory""": tracemalloc.peaked + bamb(tracemalloc.begin							),
       },     step=_lowerCamelCase,     )

  # New Code #
  # context manager to track the peak memory usage during the evaluation
  with TorchTracemalloc() as tracemalloc:
   model.eval()
   for step, batch in enumerate(_lowerCamelCase							):
    # We could avoid this line since we set the accelerator with `device_placement=True`.
    batch.to(accelerator.device							)
    with torch.no_grad():
     lowerCamelCase		:     str          =				model(**_lowerCamelCase							)
    lowerCamelCase		:     List[Any]          =				outputs.logits.argmax(dim=-1							)
    lowerCamelCase		:     str          =				accelerator.gather_for_metrics((predictions, batch["""labels"""])							)
    metric.add_batch(
        predictions=_lowerCamelCase,     references=_lowerCamelCase,     )

   lowerCamelCase		:     Optional[int]          =				metric.compute()
   # Use accelerator.print to print only on the main process.
   accelerator.print(F'''epoch {epoch}:''',     _lowerCamelCase							)
   if args.with_tracking:
    accelerator.log(
        {
            """accuracy""": eval_metric["""accuracy"""],
            """f1""": eval_metric["""f1"""],
            """train_loss""": total_loss.item() / len(_lowerCamelCase							),
        },     step=_lowerCamelCase,     )

   if checkpointing_steps == "epoch":
    lowerCamelCase		:     Any          =				F'''epoch_{epoch}'''
    if args.output_dir is not None:
     lowerCamelCase		:     Tuple          =				os.path.join(args.output_dir,     _lowerCamelCase							)
    accelerator.save_state(_lowerCamelCase							)
    # New Code #
    # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
  accelerator.print("""Memory before entering the eval : {}""".format(bamb(tracemalloc.begin							)							)							)
  accelerator.print("""Memory consumed at the end of the eval (end-begin): {}""".format(tracemalloc.used							)							)
  accelerator.print("""Peak Memory consumed during the eval (max-begin): {}""".format(tracemalloc.peaked							)							)
  accelerator.print(
      """Total Peak Memory consumed during the eval (max): {}""".format(tracemalloc.peaked + bamb(tracemalloc.begin							)							)							)
  # Logging the peak memory usage of the GPU to the tracker
  if args.with_tracking:
   accelerator.log(
       {
           """eval_total_peak_memory""": tracemalloc.peaked + bamb(tracemalloc.begin							),
       },     step=_lowerCamelCase,     )

 if args.with_tracking:
  accelerator.end_training()





def    _a (						):
 lowerCamelCase		:     Any          =				argparse.ArgumentParser(description="""Simple example of training script."""							)
 parser.add_argument(
     """--mixed_precision""",     type=_lowerCamelCase,     default=_lowerCamelCase,     choices=["""no""", """fp16""", """bf16""", """fp8"""],     help="""Whether to use mixed precision. Choose"""
     """between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."""
     """and an Nvidia Ampere GPU.""",     )
 parser.add_argument("""--cpu""",     action="""store_true""",     help="""If passed, will train on the CPU."""							)
 parser.add_argument(
     """--checkpointing_steps""",     type=_lowerCamelCase,     default=_lowerCamelCase,     help="""Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.""",     )
 parser.add_argument(
     """--resume_from_checkpoint""",     type=_lowerCamelCase,     default=_lowerCamelCase,     help="""If the training should continue from a checkpoint folder.""",     )
 parser.add_argument(
     """--with_tracking""",     action="""store_true""",     help="""Whether to load in all available experiment trackers from the environment and use them for logging.""",     )
 parser.add_argument(
     """--output_dir""",     type=_lowerCamelCase,     default=""".""",     help="""Optional save directory where all checkpoint folders will be stored. Default is the current working directory.""",     )
 parser.add_argument(
     """--logging_dir""",     type=_lowerCamelCase,     default="""logs""",     help="""Location on where to store experiment tracking logs`""",     )
 parser.add_argument(
     """--model_name_or_path""",     type=_lowerCamelCase,     help="""Path to pretrained model or model identifier from huggingface.co/models.""",     required=_lowerCamelCase,     )
 lowerCamelCase		:     int          =				parser.parse_args()
 lowerCamelCase		:     Optional[Any]          =				{"""lr""": 2e-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 16}
 training_function(_lowerCamelCase,     _lowerCamelCase							)


if __name__ == "__main__":
      main()






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,     lowerCamelCase,     lowerCamelCase							):
 # Initialise PyTorch model
 lowerCamelCase		:     List[str]          =				FNetConfig.from_json_file(lowerCamelCase							)
 print(F'''Building PyTorch model from configuration: {config}'''							)
 lowerCamelCase		:     Any          =				FNetForPreTraining(lowerCamelCase							)

 lowerCamelCase		:     List[Any]          =				restore_checkpoint(lowerCamelCase,     lowerCamelCase							)
 lowerCamelCase		:     List[str]          =				checkpoint_dict["""target"""]

 # Embeddings
 # Position IDs
 lowerCamelCase		:     Optional[Any]          =				fnet_pretraining_model.state_dict()

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

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

  lowerCamelCase		:     int          =				torch.tensor(
      pretrained_model_params["""encoder"""][F'''feed_forward_{layer}''']["""intermediate"""]["""kernel"""]							).T
  lowerCamelCase		:     List[str]          =				torch.tensor(
      pretrained_model_params["""encoder"""][F'''feed_forward_{layer}''']["""intermediate"""]["""bias"""]							)

  lowerCamelCase		:     List[Any]          =				torch.tensor(
      pretrained_model_params["""encoder"""][F'''feed_forward_{layer}''']["""output"""]["""kernel"""]							).T
  lowerCamelCase		:     str          =				torch.tensor(
      pretrained_model_params["""encoder"""][F'''feed_forward_{layer}''']["""output"""]["""bias"""]							)

  lowerCamelCase		:     Tuple          =				torch.tensor(
      pretrained_model_params["""encoder"""][F'''encoder_{layer}''']["""output_layer_norm"""]["""scale"""]							)
  lowerCamelCase		:     Optional[int]          =				torch.tensor(
      pretrained_model_params["""encoder"""][F'''encoder_{layer}''']["""output_layer_norm"""]["""bias"""]							)

 # Pooler Layers
 lowerCamelCase		:     Tuple          =				torch.tensor(pretrained_model_params["""encoder"""]["""pooler"""]["""kernel"""]							).T
 lowerCamelCase		:     List[Any]          =				torch.tensor(pretrained_model_params["""encoder"""]["""pooler"""]["""bias"""]							)

 # Masked LM Layers
 lowerCamelCase		:     Optional[Any]          =				torch.tensor(
     pretrained_model_params["""predictions_dense"""]["""kernel"""]							).T
 lowerCamelCase		:     Tuple          =				torch.tensor(
     pretrained_model_params["""predictions_dense"""]["""bias"""]							)
 lowerCamelCase		:     Tuple          =				torch.tensor(
     pretrained_model_params["""predictions_layer_norm"""]["""scale"""]							)
 lowerCamelCase		:     int          =				torch.tensor(
     pretrained_model_params["""predictions_layer_norm"""]["""bias"""]							)
 lowerCamelCase		:     Dict          =				torch.tensor(
     pretrained_model_params["""encoder"""]["""embedder"""]["""word"""]["""embedding"""]							)
 lowerCamelCase		:     Dict          =				torch.tensor(
     pretrained_model_params["""predictions_output"""]["""output_bias"""]							)
 lowerCamelCase		:     Any          =				torch.tensor(pretrained_model_params["""predictions_output"""]["""output_bias"""]							)

 # Seq Relationship Layers
 lowerCamelCase		:     int          =				torch.tensor(
     pretrained_model_params["""classification"""]["""output_kernel"""]							)
 lowerCamelCase		:     Optional[Any]          =				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__":
      _lowerCamelCase        =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.""")
      _lowerCamelCase        =parser.parse_args()
      convert_flax_checkpoint_to_pytorch(args.flax_checkpoint_path, args.fnet_config_file, args.save_path)




