import json
import os
from dataclasses import dataclass
from functools import partial
from typing import Callable

import flax.linen as nn
import jax
import jax.numpy as jnp
import joblib
import optax
import wandb
from flax import jax_utils, struct, traverse_util
from flax.serialization import from_bytes, to_bytes
from flax.training import train_state
from flax.training.common_utils import shard
from tqdm.auto import tqdm

from transformers import BigBirdConfig, FlaxBigBirdForQuestionAnswering
from transformers.models.big_bird.modeling_flax_big_bird import FlaxBigBirdForQuestionAnsweringModule



class  UpperCamelCase				( __lowercase ):

					'''simple docstring'''

					A_										=  42
					A_										=  jnp.floataa
					A_										=  True



					def 		UpperCamelCase_						(					self						) ->				Tuple:






											"""simple docstring"""
											super().setup()
											_lowerCamelCase			    =				nn.Dense(5				,				dtype=self.dtype						)





					def __call__(					self				,				*A_				,				**A_						) ->				int:






											"""simple docstring"""
											_lowerCamelCase			    =				super().__call__(*A_				,				**A_						)
											_lowerCamelCase			    =				self.cls(outputs[2]						)
											return outputs[:2] + (cls_out,)





class  UpperCamelCase				( __lowercase ):

					'''simple docstring'''

					A_										=  FlaxBigBirdForNaturalQuestionsModule





def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)						->    Any:



						'''simple docstring'''

						def cross_entropy(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase=None		):
												_lowerCamelCase			    =				logits.shape[-1]
												_lowerCamelCase			    =				(labels[..., None] == jnp.arange(__UpperCAmelCase		)[None]).astype('''f4'''		)
												_lowerCamelCase			    =				jax.nn.log_softmax(__UpperCAmelCase ,						axis=-1		)
												_lowerCamelCase			    =				-jnp.sum(labels * logits ,						axis=-1		)
												if reduction is not None:
																		_lowerCamelCase			    =				reduction(__UpperCAmelCase		)
												return loss

						_lowerCamelCase			    =				partial(__UpperCAmelCase ,						reduction=jnp.mean		)
						_lowerCamelCase			    =				cross_entropy(__UpperCAmelCase ,						__UpperCAmelCase		)
						_lowerCamelCase			    =				cross_entropy(__UpperCAmelCase ,						__UpperCAmelCase		)
						_lowerCamelCase			    =				cross_entropy(__UpperCAmelCase ,						__UpperCAmelCase		)
						return (start_loss + end_loss + pooled_loss) / 3



@dataclass
class  UpperCamelCase				:

					'''simple docstring'''

					A_										=  'google/bigbird-roberta-base'
					A_										=  3_000
					A_										=  10_500

					A_										=  128
					A_										=  3

					A_										=  1
					A_										=  5

					# tx_args
					A_										=  3E-5
					A_										=  0.0
					A_										=  20_000
					A_										=  0.00_95

					A_										=  'bigbird-roberta-natural-questions'
					A_										=  'training-expt'
					A_										=  'data/nq-training.jsonl'
					A_										=  'data/nq-validation.jsonl'



					def 		UpperCamelCase_						(					self						) ->				Dict:






											"""simple docstring"""
											os.makedirs(self.base_dir				,				exist_ok=A_						)
											_lowerCamelCase			    =				os.path.join(self.base_dir				,				self.save_dir						)
											_lowerCamelCase			    =				self.batch_size_per_device * jax.device_count()





@dataclass
class  UpperCamelCase				:

					'''simple docstring'''

					A_										=  42
					A_										=  4_096  # no dynamic padding on TPUs



					def __call__(					self				,				A_						) ->				Any:






											"""simple docstring"""
											_lowerCamelCase			    =				self.collate_fn(A_						)
											_lowerCamelCase			    =				jax.tree_util.tree_map(A_				,				A_						)
											return batch



					def 		UpperCamelCase_						(					self				,				A_						) ->				Optional[int]:






											"""simple docstring"""
											_lowerCamelCase    ,					_lowerCamelCase			    =				self.fetch_inputs(features['''input_ids''']						)
											_lowerCamelCase			    =				{
											    '''input_ids''': jnp.array(A_				,				dtype=jnp.intaa						),
											    '''attention_mask''': jnp.array(A_				,				dtype=jnp.intaa						),
											    '''start_labels''': jnp.array(features['''start_token''']				,				dtype=jnp.intaa						),
											    '''end_labels''': jnp.array(features['''end_token''']				,				dtype=jnp.intaa						),
											    '''pooled_labels''': jnp.array(features['''category''']				,				dtype=jnp.intaa						),
											}
											return batch



					def 		UpperCamelCase_						(					self				,				A_						) ->				Optional[int]:






											"""simple docstring"""
											_lowerCamelCase			    =				[self._fetch_inputs(A_						) for ids in input_ids]
											return zip(*A_						)





					def 		UpperCamelCase_						(					self				,				A_						) ->				Any:






											"""simple docstring"""
											_lowerCamelCase			    =				[1 for _ in range(len(A_						)						)]
											while len(A_						) < self.max_length:
																	input_ids.append(self.pad_id						)
																	attention_mask.append(0						)
											return input_ids, attention_mask





def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase=None		)						->    Union[str, Any]:



						'''simple docstring'''

						if seed is not None:
												_lowerCamelCase			    =				dataset.shuffle(seed=__UpperCAmelCase		)
						for i in range(len(__UpperCAmelCase		) // batch_size		):
												_lowerCamelCase			    =				dataset[i * batch_size : (i + 1) * batch_size]
												yield dict(__UpperCAmelCase		)
@partial(jax.pmap ,						axis_name='''batch'''		)
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						**__UpperCAmelCase		)						->    Any:



						'''simple docstring'''

						def loss_fn(__UpperCAmelCase		):
												_lowerCamelCase			    =				model_inputs.pop('''start_labels'''		)
												_lowerCamelCase			    =				model_inputs.pop('''end_labels'''		)
												_lowerCamelCase			    =				model_inputs.pop('''pooled_labels'''		)

												_lowerCamelCase			    =				state.apply_fn(**__UpperCAmelCase ,						params=__UpperCAmelCase ,						dropout_rng=__UpperCAmelCase ,						train=__UpperCAmelCase		)
												_lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase			    =				outputs

												return state.loss_fn(
												    __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						)

						_lowerCamelCase    ,					_lowerCamelCase			    =				jax.random.split(__UpperCAmelCase		)
						_lowerCamelCase			    =				jax.value_and_grad(__UpperCAmelCase		)
						_lowerCamelCase    ,					_lowerCamelCase			    =				grad_fn(state.params		)
						_lowerCamelCase			    =				jax.lax.pmean({'''loss''': loss} ,						axis_name='''batch'''		)
						_lowerCamelCase			    =				jax.lax.pmean(__UpperCAmelCase ,						'''batch'''		)

						_lowerCamelCase			    =				state.apply_gradients(grads=__UpperCAmelCase		)
						return state, metrics, new_drp_rng
@partial(jax.pmap ,						axis_name='''batch'''		)
def       __magic_name__(  __UpperCAmelCase ,						**__UpperCAmelCase		)						->    List[Any]:



						'''simple docstring'''

						_lowerCamelCase			    =				model_inputs.pop('''start_labels'''		)
						_lowerCamelCase			    =				model_inputs.pop('''end_labels'''		)
						_lowerCamelCase			    =				model_inputs.pop('''pooled_labels'''		)

						_lowerCamelCase			    =				state.apply_fn(**__UpperCAmelCase ,						params=state.params ,						train=__UpperCAmelCase		)
						_lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase			    =				outputs

						_lowerCamelCase			    =				state.loss_fn(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)
						_lowerCamelCase			    =				jax.lax.pmean({'''loss''': loss} ,						axis_name='''batch'''		)
						return metrics



class  UpperCamelCase				( train_state.TrainState ):

					'''simple docstring'''

					A_										=  struct.field(pytree_node=__lowercase )





@dataclass
class  UpperCamelCase				:

					'''simple docstring'''

					A_										=  42
					A_										=  42
					A_										=  42
					A_										=  42
					A_										=  42
					A_										=  42
					A_										=  None



					def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_				,				A_=None						) ->				List[str]:






											"""simple docstring"""
											_lowerCamelCase			    =				model.params
											_lowerCamelCase			    =				TrainState.create(
											    apply_fn=model.__call__				,				params=A_				,				tx=A_				,				loss_fn=A_				,				)
											if ckpt_dir is not None:
																	_lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase			    =				restore_checkpoint(A_				,				A_						)
																	_lowerCamelCase			    =				{
																	    '''lr''': args.lr,
																	    '''init_lr''': args.init_lr,
																	    '''warmup_steps''': args.warmup_steps,
																	    '''num_train_steps''': num_train_steps,
																	    '''weight_decay''': args.weight_decay,
																	}
																	_lowerCamelCase    ,					_lowerCamelCase			    =				build_tx(**A_						)
																	_lowerCamelCase			    =				train_state.TrainState(
																	    step=A_				,				apply_fn=model.__call__				,				params=A_				,				tx=A_				,				opt_state=A_				,				)
																	_lowerCamelCase			    =				args
																	_lowerCamelCase			    =				data_collator
																	_lowerCamelCase			    =				lr
																	_lowerCamelCase			    =				params
											_lowerCamelCase			    =				jax_utils.replicate(A_						)
											return state



					def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				Tuple:






											"""simple docstring"""
											_lowerCamelCase			    =				self.args
											_lowerCamelCase			    =				len(A_						) // args.batch_size

											_lowerCamelCase			    =				jax.random.PRNGKey(0						)
											_lowerCamelCase			    =				jax.random.split(A_				,				jax.device_count()						)
											for epoch in range(args.max_epochs						):
																	_lowerCamelCase			    =				jnp.array(0				,				dtype=jnp.floataa						)
																	_lowerCamelCase			    =				get_batched_dataset(A_				,				args.batch_size				,				seed=A_						)
																	_lowerCamelCase			    =				0
																	for batch in tqdm(A_				,				total=A_				,				desc=F'Running EPOCH-{epoch}'						):
																							_lowerCamelCase			    =				self.data_collator(A_						)
																							_lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase			    =				self.train_step_fn(A_				,				A_				,				**A_						)
																							running_loss += jax_utils.unreplicate(metrics['''loss''']						)
																							i += 1
																							if i % args.logging_steps == 0:
																													_lowerCamelCase			    =				jax_utils.unreplicate(state.step						)
																													_lowerCamelCase			    =				running_loss.item() / i
																													_lowerCamelCase			    =				self.scheduler_fn(state_step - 1						)

																													_lowerCamelCase			    =				self.evaluate(A_				,				A_						)
																													_lowerCamelCase			    =				{
																													    '''step''': state_step.item(),
																													    '''eval_loss''': eval_loss.item(),
																													    '''tr_loss''': tr_loss,
																													    '''lr''': lr.item(),
																													}
																													tqdm.write(str(A_						)						)
																													self.logger.log(A_				,				commit=A_						)

																							if i % args.save_steps == 0:
																													self.save_checkpoint(args.save_dir + F'-e{epoch}-s{i}'				,				state=A_						)



					def 		UpperCamelCase_						(					self				,				A_				,				A_						) ->				Optional[int]:






											"""simple docstring"""
											_lowerCamelCase			    =				get_batched_dataset(A_				,				self.args.batch_size						)
											_lowerCamelCase			    =				len(A_						) // self.args.batch_size
											_lowerCamelCase			    =				jnp.array(0				,				dtype=jnp.floataa						)
											_lowerCamelCase			    =				0
											for batch in tqdm(A_				,				total=A_				,				desc='''Evaluating ... '''						):
																	_lowerCamelCase			    =				self.data_collator(A_						)
																	_lowerCamelCase			    =				self.val_step_fn(A_				,				**A_						)
																	running_loss += jax_utils.unreplicate(metrics['''loss''']						)
																	i += 1
											return running_loss / i





					def 		UpperCamelCase_						(					self				,				A_				,				A_						) ->				Optional[int]:






											"""simple docstring"""
											_lowerCamelCase			    =				jax_utils.unreplicate(A_						)
											print(F'SAVING CHECKPOINT IN {save_dir}'				,				end=''' ... '''						)
											self.model_save_fn(A_				,				params=state.params						)
											with open(os.path.join(A_				,				'''opt_state.msgpack'''						)				,				'''wb'''						) as f:
																	f.write(to_bytes(state.opt_state						)						)
											joblib.dump(self.args				,				os.path.join(A_				,				'''args.joblib'''						)						)
											joblib.dump(self.data_collator				,				os.path.join(A_				,				'''data_collator.joblib'''						)						)
											with open(os.path.join(A_				,				'''training_state.json'''						)				,				'''w'''						) as f:
																	json.dump({'''step''': state.step.item()}				,				A_						)
											print('''DONE'''						)





def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase		)						->    List[Any]:



						'''simple docstring'''

						print(F'RESTORING CHECKPOINT FROM {save_dir}' ,						end=''' ... '''		)
						with open(os.path.join(__UpperCAmelCase ,						'''flax_model.msgpack'''		) ,						'''rb'''		) as f:
												_lowerCamelCase			    =				from_bytes(state.params ,						f.read()		)

						with open(os.path.join(__UpperCAmelCase ,						'''opt_state.msgpack'''		) ,						'''rb'''		) as f:
												_lowerCamelCase			    =				from_bytes(state.opt_state ,						f.read()		)

						_lowerCamelCase			    =				joblib.load(os.path.join(__UpperCAmelCase ,						'''args.joblib'''		)		)
						_lowerCamelCase			    =				joblib.load(os.path.join(__UpperCAmelCase ,						'''data_collator.joblib'''		)		)

						with open(os.path.join(__UpperCAmelCase ,						'''training_state.json'''		) ,						'''r'''		) as f:
												_lowerCamelCase			    =				json.load(__UpperCAmelCase		)
						_lowerCamelCase			    =				training_state['''step''']

						print('''DONE'''		)
						return params, opt_state, step, args, data_collator
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)						->    Union[str, Any]:



						'''simple docstring'''

						_lowerCamelCase			    =				num_train_steps - warmup_steps
						_lowerCamelCase			    =				optax.linear_schedule(init_value=__UpperCAmelCase ,						end_value=__UpperCAmelCase ,						transition_steps=__UpperCAmelCase		)
						_lowerCamelCase			    =				optax.linear_schedule(init_value=__UpperCAmelCase ,						end_value=1E-7 ,						transition_steps=__UpperCAmelCase		)
						_lowerCamelCase			    =				optax.join_schedules(schedules=[warmup_fn, decay_fn] ,						boundaries=[warmup_steps]		)
						return lr




def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)						->    Tuple:



						'''simple docstring'''

						def weight_decay_mask(__UpperCAmelCase		):
												_lowerCamelCase			    =				traverse_util.flatten_dict(__UpperCAmelCase		)
												_lowerCamelCase			    =				{k: (v[-1] != '''bias''' and v[-2:] != ('''LayerNorm''', '''scale''')) for k, v in params.items()}
												return traverse_util.unflatten_dict(__UpperCAmelCase		)

						_lowerCamelCase			    =				scheduler_fn(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)

						_lowerCamelCase			    =				optax.adamw(learning_rate=__UpperCAmelCase ,						weight_decay=__UpperCAmelCase ,						mask=__UpperCAmelCase		)
						return tx, lr
import argparse
import json
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from typing import List

import timm
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from torch import Tensor

from transformers import AutoImageProcessor, ResNetConfig, ResNetForImageClassification
from transformers.utils import logging


logging.set_verbosity_info()
snake_case__				= logging.get_logger()



@dataclass
class  UpperCamelCase				:

     '''simple docstring'''

     A_										=  42
     A_										=  field(default_factory=__lowercase )
     A_										=  field(default_factory=__lowercase )



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				Any:






           """simple docstring"""
           _lowerCamelCase			    =				len(list(m.modules()						)						) == 1 or isinstance(A_				,				nn.Convad						) or isinstance(A_				,				nn.BatchNormad						)
           if has_not_submodules:
                 self.traced.append(A_						)



     def __call__(					self				,				A_						) ->				Tuple:






           """simple docstring"""
           for m in self.module.modules():
                 self.handles.append(m.register_forward_hook(self._forward_hook						)						)
           self.module(A_						)
           [x.remove() for x in self.handles]
           return self





     @property
     def 		UpperCamelCase_						(					self						) ->				List[str]:






           """simple docstring"""
           # check the len of the state_dict keys to see if we have learnable params
           return list(filter(lambda A_						: len(list(x.state_dict().keys()						)						) > 0				,				self.traced						)						)





@dataclass
class  UpperCamelCase				:

     '''simple docstring'''

     A_										=  42
     A_										=  42
     A_										=  0
     A_										=  field(default_factory=__lowercase )
     A_										=  field(default_factory=__lowercase )



     def __call__(					self				,				A_						) ->				List[Any]:






           """simple docstring"""
           _lowerCamelCase			    =				Tracker(self.dest						)(A_						).parametrized
           _lowerCamelCase			    =				Tracker(self.src						)(A_						).parametrized

           _lowerCamelCase			    =				list(filter(lambda A_						: type(A_						) not in self.src_skip				,				A_						)						)
           _lowerCamelCase			    =				list(filter(lambda A_						: type(A_						) not in self.dest_skip				,				A_						)						)

           if len(A_						) != len(A_						):
                 raise Exception(
                     F'Numbers of operations are different. Source module has {len(A_						)} operations while'
                     F' destination module has {len(A_						)}.'						)

           for dest_m, src_m in zip(A_				,				A_						):
                 dest_m.load_state_dict(src_m.state_dict()						)
                 if self.verbose == 1:
                       print(F'Transfered from={src_m} to={dest_m}'						)





def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase = True		)						->    Optional[int]:



      '''simple docstring'''

      print(F'Converting {name}...'		)
      with torch.no_grad():
            _lowerCamelCase			    =				timm.create_model(__UpperCAmelCase ,						pretrained=__UpperCAmelCase		).eval()
            _lowerCamelCase			    =				ResNetForImageClassification(__UpperCAmelCase		).eval()
            _lowerCamelCase			    =				ModuleTransfer(src=__UpperCAmelCase ,						dest=__UpperCAmelCase		)
            _lowerCamelCase			    =				torch.randn((1, 3, 224, 224)		)
            module_transfer(__UpperCAmelCase		)

      assert torch.allclose(from_model(__UpperCAmelCase		) ,						our_model(__UpperCAmelCase		).logits		), "The model logits don't match the original one."

      _lowerCamelCase			    =				F'resnet{"-".join(name.split("resnet"		)		)}'
      print(__UpperCAmelCase		)

      if push_to_hub:
            our_model.push_to_hub(
                repo_path_or_name=save_directory / checkpoint_name ,						commit_message='''Add model''' ,						use_temp_dir=__UpperCAmelCase ,						)

            # we can use the convnext one
            _lowerCamelCase			    =				AutoImageProcessor.from_pretrained('''facebook/convnext-base-224-22k-1k'''		)
            image_processor.push_to_hub(
                repo_path_or_name=save_directory / checkpoint_name ,						commit_message='''Add image processor''' ,						use_temp_dir=__UpperCAmelCase ,						)

            print(F'Pushed {checkpoint_name}'		)
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase = None ,						__UpperCAmelCase = True		)						->    Optional[int]:



      '''simple docstring'''

      _lowerCamelCase			    =				'''imagenet-1k-id2label.json'''
      _lowerCamelCase			    =				1000
      _lowerCamelCase			    =				(1, num_labels)

      _lowerCamelCase			    =				'''huggingface/label-files'''
      _lowerCamelCase			    =				num_labels
      _lowerCamelCase			    =				json.load(open(hf_hub_download(__UpperCAmelCase ,						__UpperCAmelCase ,						repo_type='''dataset'''		) ,						'''r'''		)		)
      _lowerCamelCase			    =				{int(__UpperCAmelCase		): v for k, v in idalabel.items()}

      _lowerCamelCase			    =				idalabel
      _lowerCamelCase			    =				{v: k for k, v in idalabel.items()}

      _lowerCamelCase			    =				partial(__UpperCAmelCase ,						num_labels=__UpperCAmelCase ,						idalabel=__UpperCAmelCase ,						labelaid=__UpperCAmelCase		)

      _lowerCamelCase			    =				{
          '''resnet18''': ImageNetPreTrainedConfig(
              depths=[2, 2, 2, 2] ,						hidden_sizes=[64, 128, 256, 512] ,						layer_type='''basic'''		),
          '''resnet26''': ImageNetPreTrainedConfig(
              depths=[2, 2, 2, 2] ,						hidden_sizes=[256, 512, 1024, 2048] ,						layer_type='''bottleneck'''		),
          '''resnet34''': ImageNetPreTrainedConfig(
              depths=[3, 4, 6, 3] ,						hidden_sizes=[64, 128, 256, 512] ,						layer_type='''basic'''		),
          '''resnet50''': ImageNetPreTrainedConfig(
              depths=[3, 4, 6, 3] ,						hidden_sizes=[256, 512, 1024, 2048] ,						layer_type='''bottleneck'''		),
          '''resnet101''': ImageNetPreTrainedConfig(
              depths=[3, 4, 23, 3] ,						hidden_sizes=[256, 512, 1024, 2048] ,						layer_type='''bottleneck'''		),
          '''resnet152''': ImageNetPreTrainedConfig(
              depths=[3, 8, 36, 3] ,						hidden_sizes=[256, 512, 1024, 2048] ,						layer_type='''bottleneck'''		),
      }

      if model_name:
            convert_weight_and_push(__UpperCAmelCase ,						names_to_config[model_name] ,						__UpperCAmelCase ,						__UpperCAmelCase		)
      else:
            for model_name, config in names_to_config.items():
                  convert_weight_and_push(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)
      return config, expected_shape


if __name__ == "__main__":
     snake_case__				= argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '--model_name',
         default=None,
         type=str,
         help=(
             'The name of the model you wish to convert, it must be one of the supported resnet* architecture,'
             ' currently: resnet18,26,34,50,101,152. If `None`, all of them will the converted.'
         ),
     )
     parser.add_argument(
         '--pytorch_dump_folder_path',
         default=None,
         type=Path,
         required=True,
         help='Path to the output PyTorch model directory.',
     )
     parser.add_argument(
         '--push_to_hub',
         default=True,
         type=bool,
         required=False,
         help='If True, push model and image processor to the hub.',
     )

     snake_case__				= parser.parse_args()
     snake_case__				= args.pytorch_dump_folder_path
     pytorch_dump_folder_path.mkdir(exist_ok=True, parents=True)
     convert_weights_and_push(pytorch_dump_folder_path, args.model_name, args.push_to_hub)