



import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import DetrConfig, DetrForObjectDetection, DetrForSegmentation, DetrImageProcessor, ResNetConfig
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase_							:      Dict      =  logging.get_logger(__name__)
def  UpperCamelCase						(       _A			:				List[str]    )->       Union[str, Any]:





      """simple docstring"""





      if "resnet-50" in model_name:
            A__							=       ResNetConfig.from_pretrained("microsoft/resnet-50"    )
      elif "resnet-101" in model_name:
            A__							=       ResNetConfig.from_pretrained("microsoft/resnet-101"    )
      else:
            raise ValueError("Model name should include either resnet50 or resnet101"    )

      A__							=       DetrConfig(use_timm_backbone=A__			,							backbone_config=A__    )

      # set label attributes
      A__							=       "panoptic" in model_name
      if is_panoptic:
            A__							=       250
      else:
            A__							=       91
            A__							=       "huggingface/label-files"
            A__							=       "coco-detection-id2label.json"
            A__							=       json.load(open(hf_hub_download(A__			,							A__			,							repo_type="dataset"    )			,							"r"    )    )
            A__							=       {int(A__    ): v for k, v in idalabel.items()}
            A__							=       idalabel
            A__							=       {v: k for k, v in idalabel.items()}

      return config, is_panoptic
def  UpperCamelCase						(       _A			:				Union[str, Any]    )->       Tuple:





      """simple docstring"""





      A__							=       []

      # stem
      # fmt: off
      rename_keys.append(("backbone.0.body.conv1.weight", "backbone.conv_encoder.model.embedder.embedder.convolution.weight")    )
      rename_keys.append(("backbone.0.body.bn1.weight", "backbone.conv_encoder.model.embedder.embedder.normalization.weight")    )
      rename_keys.append(("backbone.0.body.bn1.bias", "backbone.conv_encoder.model.embedder.embedder.normalization.bias")    )
      rename_keys.append(("backbone.0.body.bn1.running_mean", "backbone.conv_encoder.model.embedder.embedder.normalization.running_mean")    )
      rename_keys.append(("backbone.0.body.bn1.running_var", "backbone.conv_encoder.model.embedder.embedder.normalization.running_var")    )
      # stages
      for stage_idx in range(len(config.backbone_config.depths    )    ):
            for layer_idx in range(config.backbone_config.depths[stage_idx]    ):
                  # shortcut
                  if layer_idx == 0:
                        rename_keys.append(
                            (
                                f"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.0.weight""",
                                f"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.convolution.weight""",
                            )    )
                        rename_keys.append(
                            (
                                f"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.weight""",
                                f"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.weight""",
                            )    )
                        rename_keys.append(
                            (
                                f"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.bias""",
                                f"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.bias""",
                            )    )
                        rename_keys.append(
                            (
                                f"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_mean""",
                                f"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_mean""",
                            )    )
                        rename_keys.append(
                            (
                                f"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_var""",
                                f"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_var""",
                            )    )
                  # 3 convs
                  for i in range(3    ):
                        rename_keys.append(
                            (
                                f"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.conv{i+1}.weight""",
                                f"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.convolution.weight""",
                            )    )
                        rename_keys.append(
                            (
                                f"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.weight""",
                                f"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.weight""",
                            )    )
                        rename_keys.append(
                            (
                                f"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.bias""",
                                f"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.bias""",
                            )    )
                        rename_keys.append(
                            (
                                f"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_mean""",
                                f"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_mean""",
                            )    )
                        rename_keys.append(
                            (
                                f"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_var""",
                                f"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_var""",
                            )    )
    # fmt: on

      for i in range(config.encoder_layers    ):
            # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
            rename_keys.append(
                (
                    f"""transformer.encoder.layers.{i}.self_attn.out_proj.weight""",
                    f"""encoder.layers.{i}.self_attn.out_proj.weight""",
                )    )
            rename_keys.append(
                (f"""transformer.encoder.layers.{i}.self_attn.out_proj.bias""", f"""encoder.layers.{i}.self_attn.out_proj.bias""")    )
            rename_keys.append((f"""transformer.encoder.layers.{i}.linear1.weight""", f"""encoder.layers.{i}.fc1.weight""")    )
            rename_keys.append((f"""transformer.encoder.layers.{i}.linear1.bias""", f"""encoder.layers.{i}.fc1.bias""")    )
            rename_keys.append((f"""transformer.encoder.layers.{i}.linear2.weight""", f"""encoder.layers.{i}.fc2.weight""")    )
            rename_keys.append((f"""transformer.encoder.layers.{i}.linear2.bias""", f"""encoder.layers.{i}.fc2.bias""")    )
            rename_keys.append(
                (f"""transformer.encoder.layers.{i}.norm1.weight""", f"""encoder.layers.{i}.self_attn_layer_norm.weight""")    )
            rename_keys.append(
                (f"""transformer.encoder.layers.{i}.norm1.bias""", f"""encoder.layers.{i}.self_attn_layer_norm.bias""")    )
            rename_keys.append(
                (f"""transformer.encoder.layers.{i}.norm2.weight""", f"""encoder.layers.{i}.final_layer_norm.weight""")    )
            rename_keys.append((f"""transformer.encoder.layers.{i}.norm2.bias""", f"""encoder.layers.{i}.final_layer_norm.bias""")    )
            # decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
            rename_keys.append(
                (
                    f"""transformer.decoder.layers.{i}.self_attn.out_proj.weight""",
                    f"""decoder.layers.{i}.self_attn.out_proj.weight""",
                )    )
            rename_keys.append(
                (f"""transformer.decoder.layers.{i}.self_attn.out_proj.bias""", f"""decoder.layers.{i}.self_attn.out_proj.bias""")    )
            rename_keys.append(
                (
                    f"""transformer.decoder.layers.{i}.multihead_attn.out_proj.weight""",
                    f"""decoder.layers.{i}.encoder_attn.out_proj.weight""",
                )    )
            rename_keys.append(
                (
                    f"""transformer.decoder.layers.{i}.multihead_attn.out_proj.bias""",
                    f"""decoder.layers.{i}.encoder_attn.out_proj.bias""",
                )    )
            rename_keys.append((f"""transformer.decoder.layers.{i}.linear1.weight""", f"""decoder.layers.{i}.fc1.weight""")    )
            rename_keys.append((f"""transformer.decoder.layers.{i}.linear1.bias""", f"""decoder.layers.{i}.fc1.bias""")    )
            rename_keys.append((f"""transformer.decoder.layers.{i}.linear2.weight""", f"""decoder.layers.{i}.fc2.weight""")    )
            rename_keys.append((f"""transformer.decoder.layers.{i}.linear2.bias""", f"""decoder.layers.{i}.fc2.bias""")    )
            rename_keys.append(
                (f"""transformer.decoder.layers.{i}.norm1.weight""", f"""decoder.layers.{i}.self_attn_layer_norm.weight""")    )
            rename_keys.append(
                (f"""transformer.decoder.layers.{i}.norm1.bias""", f"""decoder.layers.{i}.self_attn_layer_norm.bias""")    )
            rename_keys.append(
                (f"""transformer.decoder.layers.{i}.norm2.weight""", f"""decoder.layers.{i}.encoder_attn_layer_norm.weight""")    )
            rename_keys.append(
                (f"""transformer.decoder.layers.{i}.norm2.bias""", f"""decoder.layers.{i}.encoder_attn_layer_norm.bias""")    )
            rename_keys.append(
                (f"""transformer.decoder.layers.{i}.norm3.weight""", f"""decoder.layers.{i}.final_layer_norm.weight""")    )
            rename_keys.append((f"""transformer.decoder.layers.{i}.norm3.bias""", f"""decoder.layers.{i}.final_layer_norm.bias""")    )

      # convolutional projection + query embeddings + layernorm of decoder + class and bounding box heads
      rename_keys.extend(
          [
              ("input_proj.weight", "input_projection.weight"),
              ("input_proj.bias", "input_projection.bias"),
              ("query_embed.weight", "query_position_embeddings.weight"),
              ("transformer.decoder.norm.weight", "decoder.layernorm.weight"),
              ("transformer.decoder.norm.bias", "decoder.layernorm.bias"),
              ("class_embed.weight", "class_labels_classifier.weight"),
              ("class_embed.bias", "class_labels_classifier.bias"),
              ("bbox_embed.layers.0.weight", "bbox_predictor.layers.0.weight"),
              ("bbox_embed.layers.0.bias", "bbox_predictor.layers.0.bias"),
              ("bbox_embed.layers.1.weight", "bbox_predictor.layers.1.weight"),
              ("bbox_embed.layers.1.bias", "bbox_predictor.layers.1.bias"),
              ("bbox_embed.layers.2.weight", "bbox_predictor.layers.2.weight"),
              ("bbox_embed.layers.2.bias", "bbox_predictor.layers.2.bias"),
          ]    )

      return rename_keys
def  UpperCamelCase						(       _A			:				Tuple			,							_A			:				List[str]			,							_A			:				int    )->       Optional[Any]:





      """simple docstring"""





      A__							=       state_dict.pop(A__    )
      A__							=       val
def  UpperCamelCase						(       _A			:				Any			,							_A			:				Optional[int]=False    )->       List[str]:





      """simple docstring"""





      A__							=       ""
      if is_panoptic:
            A__							=       "detr."

      # first: transformer encoder
      for i in range(6    ):
            # read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
            A__							=       state_dict.pop(f"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight"""    )
            A__							=       state_dict.pop(f"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias"""    )
            # next, add query, keys and values (in that order) to the state dict
            A__							=       in_proj_weight[:256, :]
            A__							=       in_proj_bias[:256]
            A__							=       in_proj_weight[256:512, :]
            A__							=       in_proj_bias[256:512]
            A__							=       in_proj_weight[-256:, :]
            A__							=       in_proj_bias[-256:]
      # next: transformer decoder (which is a bit more complex because it also includes cross-attention)
      for i in range(6    ):
            # read in weights + bias of input projection layer of self-attention
            A__							=       state_dict.pop(f"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight"""    )
            A__							=       state_dict.pop(f"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias"""    )
            # next, add query, keys and values (in that order) to the state dict
            A__							=       in_proj_weight[:256, :]
            A__							=       in_proj_bias[:256]
            A__							=       in_proj_weight[256:512, :]
            A__							=       in_proj_bias[256:512]
            A__							=       in_proj_weight[-256:, :]
            A__							=       in_proj_bias[-256:]
            # read in weights + bias of input projection layer of cross-attention
            A__							=       state_dict.pop(
                f"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight"""    )
            A__							=       state_dict.pop(f"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias"""    )
            # next, add query, keys and values (in that order) of cross-attention to the state dict
            A__							=       in_proj_weight_cross_attn[:256, :]
            A__							=       in_proj_bias_cross_attn[:256]
            A__							=       in_proj_weight_cross_attn[256:512, :]
            A__							=       in_proj_bias_cross_attn[256:512]
            A__							=       in_proj_weight_cross_attn[-256:, :]
            A__							=       in_proj_bias_cross_attn[-256:]
def  UpperCamelCase						(       )->       int:





      """simple docstring"""





      A__							=       "http://images.cocodataset.org/val2017/000000039769.jpg"
      A__							=       Image.open(requests.get(A__			,							stream=A__    ).raw    )

      return im
@torch.no_grad()
def  UpperCamelCase						(       _A			:				int			,							_A			:				Dict=None			,							_A			:				List[str]=False    )->       Tuple:





      """simple docstring"""





      A__						,       A__							=       get_detr_config(A__    )

      # load original model from torch hub
      A__							=       {
          "detr-resnet-50": "detr_resnet50",
          "detr-resnet-101": "detr_resnet101",
      }
      logger.info(f"""Converting model {model_name}..."""    )
      A__							=       torch.hub.load("facebookresearch/detr"			,							model_name_to_original_name[model_name]			,							pretrained=A__    ).eval()
      A__							=       detr.state_dict()
      # rename keys
      for src, dest in create_rename_keys(A__    ):
            if is_panoptic:
                  A__							=       "detr." + src
            rename_key(A__			,							A__			,							A__    )
      # query, key and value matrices need special treatment
      read_in_q_k_v(A__			,							is_panoptic=A__    )
      # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
      A__							=       "detr.model." if is_panoptic else "model."
      for key in state_dict.copy().keys():
            if is_panoptic:
                  if (
                      key.startswith("detr"    )
                      and not key.startswith("class_labels_classifier"    )
                      and not key.startswith("bbox_predictor"    )
                  ):
                        A__							=       state_dict.pop(A__    )
                        A__							=       val
                  elif "class_labels_classifier" in key or "bbox_predictor" in key:
                        A__							=       state_dict.pop(A__    )
                        A__							=       val
                  elif key.startswith("bbox_attention"    ) or key.startswith("mask_head"    ):
                        continue
                  else:
                        A__							=       state_dict.pop(A__    )
                        A__							=       val
            else:
                  if not key.startswith("class_labels_classifier"    ) and not key.startswith("bbox_predictor"    ):
                        A__							=       state_dict.pop(A__    )
                        A__							=       val

    # finally, create HuggingFace model and load state dict
      A__							=       DetrForSegmentation(A__    ) if is_panoptic else DetrForObjectDetection(A__    )
      model.load_state_dict(A__    )
      model.eval()

      # verify our conversion on an image
      A__							=       "coco_panoptic" if is_panoptic else "coco_detection"
      A__							=       DetrImageProcessor(format=A__    )

      A__							=       processor(images=prepare_img()			,							return_tensors="pt"    )
      A__							=       encoding["pixel_values"]

      A__							=       detr(A__    )
      A__							=       model(A__    )

      assert torch.allclose(outputs.logits			,							original_outputs["pred_logits"]			,							atol=1E-3    )
      assert torch.allclose(outputs.pred_boxes			,							original_outputs["pred_boxes"]			,							atol=1E-3    )
      if is_panoptic:
            assert torch.allclose(outputs.pred_masks			,							original_outputs["pred_masks"]			,							atol=1E-4    )
      print("Looks ok!"    )

      if pytorch_dump_folder_path is not None:
            # Save model and image processor
            logger.info(f"""Saving PyTorch model and image processor to {pytorch_dump_folder_path}..."""    )
            Path(A__    ).mkdir(exist_ok=A__    )
            model.save_pretrained(A__    )
            processor.save_pretrained(A__    )

      if push_to_hub:
            # Upload model and image processor to the hub
            logger.info("Uploading PyTorch model and image processor to the hub..."    )
            model.push_to_hub(f"""nielsr/{model_name}"""    )
            processor.push_to_hub(f"""nielsr/{model_name}"""    )


if __name__ == "__main__":
      UpperCAmelCase_							:      List[Any]      =  argparse.ArgumentParser()

      parser.add_argument(
          "--model_name",
          default="detr-resnet-50",
          type=str,
          choices=["detr-resnet-50", "detr-resnet-101"],
          help="Name of the DETR model you\'d like to convert.",
      )
      parser.add_argument(
          "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
      )
      parser.add_argument("--push_to_hub", action="store_true", help="Whether to push the model to the hub or not.")
      UpperCAmelCase_							:      List[str]      =  parser.parse_args()
      convert_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)










import argparse
import gc
import json
import os
import re

import torch
from huggingface_hub import hf_hub_download

from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerFast, RwkvConfig
from transformers.modeling_utils import WEIGHTS_INDEX_NAME, shard_checkpoint


UpperCAmelCase_							:      Dict      =  {
    "169M": 12,
    "430M": 24,
    "1B5": 24,
    "3B": 32,
    "7B": 32,
    "14B": 40,
}

UpperCAmelCase_							:      Any      =  {
    "169M": 768,
    "430M": 1_024,
    "1B5": 2_048,
    "3B": 2_560,
    "7B": 4_096,
    "14B": 5_120,
}
def  UpperCamelCase						(       _A			:				Dict    )->       Optional[int]:





						"""simple docstring"""





						A__							=       list(state_dict.keys()    )
						for name in state_dict_keys:
												A__							=       state_dict.pop(_A    )
												# emb -> embedding
												if name.startswith("emb."    ):
																		A__							=       name.replace("emb."			,							"embeddings."    )
												# ln_0 -> pre_ln (only present at block 0)
												if name.startswith("blocks.0.ln0"    ):
																		A__							=       name.replace("blocks.0.ln0"			,							"blocks.0.pre_ln"    )
												# att -> attention
												A__							=       re.sub(R"blocks\.(\d+)\.att"			,							R"blocks.\1.attention"			,							_A    )
												# ffn -> feed_forward
												A__							=       re.sub(R"blocks\.(\d+)\.ffn"			,							R"blocks.\1.feed_forward"			,							_A    )
												# time_mix_k -> time_mix_key and reshape
												if name.endswith(".time_mix_k"    ):
																		A__							=       name.replace(".time_mix_k"			,							".time_mix_key"    )
												# time_mix_v -> time_mix_value and reshape
												if name.endswith(".time_mix_v"    ):
																		A__							=       name.replace(".time_mix_v"			,							".time_mix_value"    )
												# time_mix_r -> time_mix_key and reshape
												if name.endswith(".time_mix_r"    ):
																		A__							=       name.replace(".time_mix_r"			,							".time_mix_receptance"    )

												if name != "head.weight":
																		A__							=       "rwkv." + name

												A__							=       weight
						return state_dict
def  UpperCamelCase						(       _A			:				str			,							_A			:				List[Any]			,							_A			:				List[Any]			,							_A			:				int=None			,							_A			:				List[str]=None			,							_A			:				Dict=False			,							_A			:				List[Any]=None    )->       str:





						"""simple docstring"""





						if tokenizer_file is None:
												print("No `--tokenizer_file` provided, we will use the default tokenizer."    )
												A__							=       50277
												A__							=       AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b"    )
						else:
												A__							=       PreTrainedTokenizerFast(tokenizer_file=_A    )
												A__							=       len(_A    )
						tokenizer.save_pretrained(_A    )

						# 2. Build the config
						A__							=       list(NUM_HIDDEN_LAYERS_MAPPING.keys()    )
						if size is None:
												# Try to infer size from the checkpoint name
												for candidate in possible_sizes:
																		if candidate in checkpoint_file:
																								A__							=       candidate
																								break
												if size is None:
																		raise ValueError("Could not infer the size, please provide it with the `--size` argument."    )
						if size not in possible_sizes:
												raise ValueError(f"""`size` should be one of {possible_sizes}, got {size}."""    )

						A__							=       RwkvConfig(
						    vocab_size=_A			,							num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size]			,							hidden_size=HIDEN_SIZE_MAPPING[size]			,							)
						config.save_pretrained(_A    )

						# 3. Download model file then convert state_dict
						A__							=       hf_hub_download(_A			,							_A    )
						A__							=       torch.load(_A			,							map_location="cpu"    )
						A__							=       convert_state_dict(_A    )

						# 4. Split in shards and save
						A__						,       A__							=       shard_checkpoint(_A    )
						for shard_file, shard in shards.items():
												torch.save(_A			,							os.path.join(_A			,							_A    )    )

						if index is not None:
												A__							=       os.path.join(_A			,							_A    )
												# Save the index as well
												with open(_A			,							"w"			,							encoding="utf-8"    ) as f:
																		A__							=       json.dumps(_A			,							indent=2			,							sort_keys=_A    ) + "\n"
																		f.write(_A    )

												# 5. Clean up shards (for some reason the file PyTorch saves take the same space as the whole state_dict
												print(
												    "Cleaning up shards. This may error with an OOM error, it this is the case don't worry you still have converted the model."    )
												A__							=       list(shards.keys()    )

												del state_dict
												del shards
												gc.collect()

												for shard_file in shard_files:
																		A__							=       torch.load(os.path.join(_A			,							_A    )    )
																		torch.save({k: v.cpu().clone() for k, v in state_dict.items()}			,							os.path.join(_A			,							_A    )    )

						del state_dict
						gc.collect()

						if push_to_hub:
												if model_name is None:
																		raise ValueError("Please provide a `model_name` to push the model to the Hub."    )
												A__							=       AutoModelForCausalLM.from_pretrained(_A    )
												model.push_to_hub(_A			,							max_shard_size="2GB"    )
												tokenizer.push_to_hub(_A    )


if __name__ == "__main__":
						UpperCAmelCase_							:      Union[str, Any]      =  argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    "--repo_id", default=None, type=str, required=True, help="Repo ID from which to pull the checkpoint."
						)
						parser.add_argument(
						    "--checkpoint_file", default=None, type=str, required=True, help="Name of the checkpoint file in the repo."
						)
						parser.add_argument(
						    "--output_dir", default=None, type=str, required=True, help="Where to save the converted model."
						)
						parser.add_argument(
						    "--tokenizer_file",
						    default=None,
						    type=str,
						    help="Path to the tokenizer file to use (if not provided, only the model is converted).",
						)
						parser.add_argument(
						    "--size",
						    default=None,
						    type=str,
						    help="Size of the model. Will be inferred from the `checkpoint_file` if not passed.",
						)
						parser.add_argument(
						    "--push_to_hub",
						    action="store_true",
						    help="Push to the Hub the converted model.",
						)
						parser.add_argument(
						    "--model_name",
						    default=None,
						    type=str,
						    help="Name of the pushed model on the Hub, including the username / organization.",
						)

						UpperCAmelCase_							:      Optional[int]      =  parser.parse_args()
						convert_rmkv_checkpoint_to_hf_format(
						    args.repo_id,
						    args.checkpoint_file,
						    args.output_dir,
						    size=args.size,
						    tokenizer_file=args.tokenizer_file,
						    push_to_hub=args.push_to_hub,
						    model_name=args.model_name,
						)





