
import argparse
import glob
import logging
import os
from argparse import Namespace
from importlib import import_module

import numpy as np
import torch
from lightning_base import BaseTransformer, add_generic_args, generic_train
from seqeval.metrics import accuracy_score, fa_score, precision_score, recall_score
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader, TensorDataset
from utils_ner import TokenClassificationTask


_UpperCAmelCase :			Optional[int]              =	logging.getLogger(__name__)


class     lowercase (		_SCREAMING_SNAKE_CASE       ):
		__lowercase	:	Union[str, Any]     					=    "token-classification"



		def __init__(			self    ,				A_	)  -> List[Any]:
				"""simple docstring"""
				if type(A_	) == dict:
						UpperCamelCase          =					Namespace(**A_	)
				UpperCamelCase          =					import_module('tasks'	)
				try:
						UpperCamelCase          =					getattr(A_    ,				hparams.task_type	)
						UpperCamelCase          =					token_classification_task_clazz()
				except AttributeError:
						raise ValueError(
						    F'''Task {hparams.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. '''
						    F'''Available tasks classes are: {TokenClassificationTask.__subclasses__()}'''	)
				UpperCamelCase          =					self.token_classification_task.get_labels(hparams.labels	)
				UpperCamelCase          =					CrossEntropyLoss().ignore_index
				super().__init__(A_    ,				len(self.labels	)    ,				self.mode	)



		def __UpperCamelCase			(			self    ,				**A_	)  -> Optional[int]:
				"""simple docstring"""
				return self.model(**A_	)



		def __UpperCamelCase			(			self    ,				A_    ,				A_	)  -> int:
				"""simple docstring"""
				UpperCamelCase          =					{'input_ids': batch[0], 'attention_mask': batch[1], 'labels': batch[3]}
				if self.config.model_type != "distilbert":
						UpperCamelCase          =					(
						    batch[2] if self.config.model_type in ['bert', 'xlnet'] else None
						)  # XLM and RoBERTa don"t use token_type_ids

				UpperCamelCase          =					self(**A_	)
				UpperCamelCase          =					outputs[0]
				# tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
				return {"loss": loss}



		def __UpperCamelCase			(			self	)  -> Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase          =					self.hparams
				for mode in ["train", "dev", "test"]:
						UpperCamelCase          =					self._feature_file(A_	)
						if os.path.exists(A_	) and not args.overwrite_cache:
								logger.info('Loading features from cached file %s'    ,				A_	)
								UpperCamelCase          =					torch.load(A_	)
						else:
								logger.info('Creating features from dataset file at %s'    ,				args.data_dir	)
								UpperCamelCase          =					self.token_classification_task.read_examples_from_file(args.data_dir    ,				A_	)
								UpperCamelCase          =					self.token_classification_task.convert_examples_to_features(
								    A_    ,				self.labels    ,				args.max_seq_length    ,				self.tokenizer    ,				cls_token_at_end=bool(self.config.model_type in ['xlnet']	)    ,				cls_token=self.tokenizer.cls_token    ,				cls_token_segment_id=2 if self.config.model_type in ['xlnet'] else 0    ,				sep_token=self.tokenizer.sep_token    ,				sep_token_extra=A_    ,				pad_on_left=bool(self.config.model_type in ['xlnet']	)    ,				pad_token=self.tokenizer.pad_token_id    ,				pad_token_segment_id=self.tokenizer.pad_token_type_id    ,				pad_token_label_id=self.pad_token_label_id    ,				)
								logger.info('Saving features into cached file %s'    ,				A_	)
								torch.save(A_    ,				A_	)



		def __UpperCamelCase			(			self    ,				A_    ,				A_    ,				A_ = False	)  -> DataLoader:
				"""simple docstring"""
				UpperCamelCase          =					self._feature_file(A_	)
				logger.info('Loading features from cached file %s'    ,				A_	)
				UpperCamelCase          =					torch.load(A_	)
				UpperCamelCase          =					torch.tensor([f.input_ids for f in features]    ,				dtype=torch.long	)
				UpperCamelCase          =					torch.tensor([f.attention_mask for f in features]    ,				dtype=torch.long	)
				if features[0].token_type_ids is not None:
						UpperCamelCase          =					torch.tensor([f.token_type_ids for f in features]    ,				dtype=torch.long	)
				else:
						UpperCamelCase          =					torch.tensor([0 for f in features]    ,				dtype=torch.long	)
						# HACK(we will not use this anymore soon)
				UpperCamelCase          =					torch.tensor([f.label_ids for f in features]    ,				dtype=torch.long	)
				return DataLoader(
				    TensorDataset(A_    ,				A_    ,				A_    ,				A_	)    ,				batch_size=A_	)



		def __UpperCamelCase			(			self    ,				A_    ,				A_	)  -> Tuple:
				"""simple docstring"""
				"""Compute validation""" ""
				UpperCamelCase          =					{'input_ids': batch[0], 'attention_mask': batch[1], 'labels': batch[3]}
				if self.config.model_type != "distilbert":
						UpperCamelCase          =					(
						    batch[2] if self.config.model_type in ['bert', 'xlnet'] else None
						)  # XLM and RoBERTa don"t use token_type_ids
				UpperCamelCase          =					self(**A_	)
				UpperCamelCase     ,  UpperCamelCase          =					outputs[:2]
				UpperCamelCase          =					logits.detach().cpu().numpy()
				UpperCamelCase          =					inputs['labels'].detach().cpu().numpy()
				return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}



		def __UpperCamelCase			(			self    ,				A_	)  -> Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase          =					torch.stack([x['val_loss'] for x in outputs]	).mean()
				UpperCamelCase          =					np.concatenate([x['pred'] for x in outputs]    ,				axis=0	)
				UpperCamelCase          =					np.argmax(A_    ,				axis=2	)
				UpperCamelCase          =					np.concatenate([x['target'] for x in outputs]    ,				axis=0	)

				UpperCamelCase          =					dict(enumerate(self.labels	)	)
				UpperCamelCase          =					[[] for _ in range(out_label_ids.shape[0]	)]
				UpperCamelCase          =					[[] for _ in range(out_label_ids.shape[0]	)]

				for i in range(out_label_ids.shape[0]	):
						for j in range(out_label_ids.shape[1]	):
								if out_label_ids[i, j] != self.pad_token_label_id:
										out_label_list[i].append(label_map[out_label_ids[i][j]]	)
										preds_list[i].append(label_map[preds[i][j]]	)

				UpperCamelCase          =					{
				    'val_loss': val_loss_mean,
				    'accuracy_score': accuracy_score(A_    ,				A_	),
				    'precision': precision_score(A_    ,				A_	),
				    'recall': recall_score(A_    ,				A_	),
				    'f1': fa_score(A_    ,				A_	),
				}

				UpperCamelCase          =					dict(results.items()	)
				UpperCamelCase          =					results
				return ret, preds_list, out_label_list



		def __UpperCamelCase			(			self    ,				A_	)  -> Dict:
				"""simple docstring"""
				UpperCamelCase     ,  UpperCamelCase     ,  UpperCamelCase          =					self._eval_end(A_	)
				UpperCamelCase          =					ret['log']
				return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}



		def __UpperCamelCase			(			self    ,				A_	)  -> Dict:
				"""simple docstring"""
				UpperCamelCase     ,  UpperCamelCase     ,  UpperCamelCase          =					self._eval_end(A_	)

				# Converting to the dict required by pl
				# https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
				# pytorch_lightning/trainer/logging.py#L139
				UpperCamelCase          =					ret['log']
				# `val_loss` is the key returned by `self._eval_end()` but actually refers to `test_loss`
				return {"avg_test_loss": logs["val_loss"], "log": logs, "progress_bar": logs}



		@staticmethod
		def __UpperCamelCase			(			A_    ,				A_	)  -> Dict:
				"""simple docstring"""
				BaseTransformer.add_model_specific_args(A_    ,				A_	)
				parser.add_argument(
				    '--task_type'    ,				default='NER'    ,				type=A_    ,				help='Task type to fine tune in training (e.g. NER, POS, etc)'	)
				parser.add_argument(
				    '--max_seq_length'    ,				default=128    ,				type=A_    ,				help=(
				        'The maximum total input sequence length after tokenization. Sequences longer '
				        'than this will be truncated, sequences shorter will be padded.'
				    )    ,				)

				parser.add_argument(
				    '--labels'    ,				default=''    ,				type=A_    ,				help='Path to a file containing all labels. If not specified, CoNLL-2003 labels are used.'    ,				)
				parser.add_argument(
				    '--gpus'    ,				default=0    ,				type=A_    ,				help='The number of GPUs allocated for this, it is by default 0 meaning none'    ,				)

				parser.add_argument(
				    '--overwrite_cache'    ,				action='store_true'    ,				help='Overwrite the cached training and evaluation sets'	)

				return parser


if __name__ == "__main__":
				_UpperCAmelCase :			Optional[int]              =	argparse.ArgumentParser()
				add_generic_args(parser, os.getcwd())
				_UpperCAmelCase :			Tuple              =	NERTransformer.add_model_specific_args(parser, os.getcwd())
				_UpperCAmelCase :			Any              =	parser.parse_args()
				_UpperCAmelCase :			Dict              =	NERTransformer(args)
				_UpperCAmelCase :			Optional[Any]              =	generic_train(model, args)

				if args.do_predict:
								# See https://github.com/huggingface/transformers/issues/3159
								# pl use this default format to create a checkpoint:
								# https://github.com/PyTorchLightning/pytorch-lightning/blob/master\
								# /pytorch_lightning/callbacks/model_checkpoint.py#L322
								_UpperCAmelCase :			Tuple              =	sorted(glob.glob(os.path.join(args.output_dir, "checkpoint-epoch=*.ckpt"), recursive=True))
								_UpperCAmelCase :			Tuple              =	model.load_from_checkpoint(checkpoints[-1])
								trainer.test(model)


import argparse

import requests
import torch

# pip3 install salesforce-lavis
# I'm actually installing a slightly modified version: pip3 install git+https://github.com/nielsrogge/LAVIS.git@fix_lavis
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipaConfig,
    BlipaForConditionalGeneration,
    BlipaProcessor,
    BlipaVisionConfig,
    BlipImageProcessor,
    OPTConfig,
    TaConfig,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD


def 				A	( )       ->  int:
		'''simple docstring'''

		UpperCamelCase          =					'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
		UpperCamelCase          =					Image.open(requests.get(lowercase   ,	stream=lowercase			).raw			).convert('RGB'			)

		return image


def 				A	( lowercase			)       ->  Any:
		'''simple docstring'''

		UpperCamelCase          =					[]
		# fmt: off

		# vision encoder
		rename_keys.append(('visual_encoder.cls_token', 'vision_model.embeddings.class_embedding')			)
		rename_keys.append(('visual_encoder.pos_embed', 'vision_model.embeddings.position_embedding')			)
		rename_keys.append(('visual_encoder.patch_embed.proj.weight', 'vision_model.embeddings.patch_embedding.weight')			)
		rename_keys.append(('visual_encoder.patch_embed.proj.bias', 'vision_model.embeddings.patch_embedding.bias')			)
		rename_keys.append(('ln_vision.weight', 'vision_model.post_layernorm.weight')			)
		rename_keys.append(('ln_vision.bias', 'vision_model.post_layernorm.bias')			)

		for i in range(config.vision_config.num_hidden_layers			):
				rename_keys.append((f'''visual_encoder.blocks.{i}.norm1.weight''', f'''vision_model.encoder.layers.{i}.layer_norm1.weight''')			)
				rename_keys.append((f'''visual_encoder.blocks.{i}.norm1.bias''', f'''vision_model.encoder.layers.{i}.layer_norm1.bias''')			)
				rename_keys.append((f'''visual_encoder.blocks.{i}.norm2.weight''', f'''vision_model.encoder.layers.{i}.layer_norm2.weight''')			)
				rename_keys.append((f'''visual_encoder.blocks.{i}.norm2.bias''', f'''vision_model.encoder.layers.{i}.layer_norm2.bias''')			)
				rename_keys.append((f'''visual_encoder.blocks.{i}.attn.qkv.weight''', f'''vision_model.encoder.layers.{i}.self_attn.qkv.weight''')			)
				rename_keys.append((f'''visual_encoder.blocks.{i}.attn.proj.weight''', f'''vision_model.encoder.layers.{i}.self_attn.projection.weight''',)			)
				rename_keys.append((f'''visual_encoder.blocks.{i}.attn.proj.bias''', f'''vision_model.encoder.layers.{i}.self_attn.projection.bias''')			)
				rename_keys.append((f'''visual_encoder.blocks.{i}.mlp.fc1.weight''', f'''vision_model.encoder.layers.{i}.mlp.fc1.weight''')			)
				rename_keys.append((f'''visual_encoder.blocks.{i}.mlp.fc1.bias''', f'''vision_model.encoder.layers.{i}.mlp.fc1.bias''')			)
				rename_keys.append((f'''visual_encoder.blocks.{i}.mlp.fc2.weight''', f'''vision_model.encoder.layers.{i}.mlp.fc2.weight''')			)
				rename_keys.append((f'''visual_encoder.blocks.{i}.mlp.fc2.bias''', f'''vision_model.encoder.layers.{i}.mlp.fc2.bias''')			)

		# QFormer
		rename_keys.append(('Qformer.bert.embeddings.LayerNorm.weight', 'qformer.layernorm.weight')			)
		rename_keys.append(('Qformer.bert.embeddings.LayerNorm.bias', 'qformer.layernorm.bias')			)

		# fmt: on
		return rename_keys


def 				A	( lowercase   ,	lowercase   ,	lowercase			)       ->  Dict:
		'''simple docstring'''

		UpperCamelCase          =					dct.pop(lowercase			)
		UpperCamelCase          =					val


def 				A	( lowercase   ,	lowercase			)       ->  List[str]:
		'''simple docstring'''

		for i in range(config.vision_config.num_hidden_layers			):
				# read in original q and v biases
				UpperCamelCase          =					state_dict.pop(f'''visual_encoder.blocks.{i}.attn.q_bias'''			)
				UpperCamelCase          =					state_dict.pop(f'''visual_encoder.blocks.{i}.attn.v_bias'''			)

				# next, set bias in the state dict
				UpperCamelCase          =					torch.cat((q_bias, torch.zeros_like(lowercase   ,	requires_grad=lowercase			), v_bias)			)
				UpperCamelCase          =					qkv_bias


def 				A	( lowercase   ,	lowercase			)       ->  int:
		'''simple docstring'''

		UpperCamelCase          =					364 if 'coco' in model_name else 224
		UpperCamelCase          =					BlipaVisionConfig(image_size=lowercase			).to_dict()

		# make sure the models have proper bos_token_id and eos_token_id set (important for generation)
		# seems like flan-T5 models don't have bos_token_id properly set?
		if "opt-2.7b" in model_name:
				UpperCamelCase          =					OPTConfig.from_pretrained('facebook/opt-2.7b'   ,	eos_token_id=lowercase			).to_dict()
		elif "opt-6.7b" in model_name:
				UpperCamelCase          =					OPTConfig.from_pretrained('facebook/opt-6.7b'   ,	eos_token_id=lowercase			).to_dict()
		elif "t5-xl" in model_name:
				UpperCamelCase          =					TaConfig.from_pretrained('google/flan-t5-xl'   ,	dense_act_fn='gelu'   ,	bos_token_id=1			).to_dict()
		elif "t5-xxl" in model_name:
				UpperCamelCase          =					TaConfig.from_pretrained('google/flan-t5-xxl'   ,	dense_act_fn='gelu'   ,	bos_token_id=1			).to_dict()

		UpperCamelCase          =					BlipaConfig(vision_config=lowercase   ,	text_config=lowercase			)

		return config, image_size


@torch.no_grad()
def 				A	( lowercase   ,	lowercase=None   ,	lowercase=False			)       ->  Union[str, Any]:
		'''simple docstring'''

		UpperCamelCase          =					(
		    AutoTokenizer.from_pretrained('facebook/opt-2.7b'			)
		    if 'opt' in model_name
		    else AutoTokenizer.from_pretrained('google/flan-t5-xl'			)
		)
		UpperCamelCase          =					tokenizer('\n'   ,	add_special_tokens=lowercase			).input_ids[0]
		UpperCamelCase     ,  UpperCamelCase          =					get_blipa_config(lowercase   ,	eos_token_id=lowercase			)

		UpperCamelCase          =					BlipaForConditionalGeneration(lowercase			).eval()

		UpperCamelCase          =					{
		    'blip2-opt-2.7b': ('blip2_opt', 'pretrain_opt2.7b'),
		    'blip2-opt-6.7b': ('blip2_opt', 'pretrain_opt6.7b'),
		    'blip2-opt-2.7b-coco': ('blip2_opt', 'caption_coco_opt2.7b'),
		    'blip2-opt-6.7b-coco': ('blip2_opt', 'caption_coco_opt6.7b'),
		    'blip2-flan-t5-xl': ('blip2_t5', 'pretrain_flant5xl'),
		    'blip2-flan-t5-xl-coco': ('blip2_t5', 'caption_coco_flant5xl'),
		    'blip2-flan-t5-xxl': ('blip2_t5', 'pretrain_flant5xxl'),
		}

		UpperCamelCase     ,  UpperCamelCase          =					model_name_to_original[model_name]

		# load original model
		print('Loading original model...'			)
		UpperCamelCase          =					'cuda' if torch.cuda.is_available() else 'cpu'
		UpperCamelCase     ,  UpperCamelCase     ,  UpperCamelCase          =					load_model_and_preprocess(
		    name=lowercase   ,	model_type=lowercase   ,	is_eval=lowercase   ,	device=lowercase			)
		original_model.eval()
		print('Done!'			)

		# update state dict keys
		UpperCamelCase          =					original_model.state_dict()
		UpperCamelCase          =					create_rename_keys(lowercase			)
		for src, dest in rename_keys:
				rename_key(lowercase   ,	lowercase   ,	lowercase			)

		# some keys can be renamed efficiently
		for key, val in state_dict.copy().items():
				UpperCamelCase          =					state_dict.pop(lowercase			)
				if key.startswith('Qformer.bert'			):
						UpperCamelCase          =					key.replace('Qformer.bert'   ,	'qformer'			)
				if "attention.self" in key:
						UpperCamelCase          =					key.replace('self'   ,	'attention'			)
				if "opt_proj" in key:
						UpperCamelCase          =					key.replace('opt_proj'   ,	'language_projection'			)
				if "t5_proj" in key:
						UpperCamelCase          =					key.replace('t5_proj'   ,	'language_projection'			)
				if key.startswith('opt'			):
						UpperCamelCase          =					key.replace('opt'   ,	'language'			)
				if key.startswith('t5'			):
						UpperCamelCase          =					key.replace('t5'   ,	'language'			)
				UpperCamelCase          =					val

		# read in qv biases
		read_in_q_v_bias(lowercase   ,	lowercase			)

		UpperCamelCase     ,  UpperCamelCase          =					hf_model.load_state_dict(lowercase   ,	strict=lowercase			)
		assert len(lowercase			) == 0
		assert unexpected_keys == ["qformer.embeddings.position_ids"]

		UpperCamelCase          =					load_demo_image()
		UpperCamelCase          =					vis_processors['eval'](lowercase			).unsqueeze(0			).to(lowercase			)
		UpperCamelCase          =					tokenizer(['\n']   ,	return_tensors='pt'			).input_ids.to(lowercase			)

		# create processor
		UpperCamelCase          =					BlipImageProcessor(
		    size={'height': image_size, 'width': image_size}   ,	image_mean=lowercase   ,	image_std=lowercase			)
		UpperCamelCase          =					BlipaProcessor(image_processor=lowercase   ,	tokenizer=lowercase			)
		UpperCamelCase          =					processor(images=lowercase   ,	return_tensors='pt'			).pixel_values.to(lowercase			)

		# make sure processor creates exact same pixel values
		assert torch.allclose(lowercase   ,	lowercase			)

		original_model.to(lowercase			)
		hf_model.to(lowercase			)
		with torch.no_grad():
				if "opt" in model_name:
						UpperCamelCase          =					original_model({'image': original_pixel_values, 'text_input': ['']}			).logits
						UpperCamelCase          =					hf_model(lowercase   ,	lowercase			).logits
				else:
						UpperCamelCase          =					original_model(
						    {'image': original_pixel_values, 'text_input': ['\n'], 'text_output': ['\n']}			).logits
						UpperCamelCase          =					input_ids.masked_fill(input_ids == tokenizer.pad_token_id   ,	-100			)
						UpperCamelCase          =					hf_model(lowercase   ,	lowercase   ,	labels=lowercase			).logits

		assert original_logits.shape == logits.shape
		print('First values of original logits:'   ,	original_logits[0, :3, :3]			)
		print('First values of HF logits:'   ,	logits[0, :3, :3]			)

		# assert values
		if model_name == "blip2-flan-t5-xl":
				UpperCamelCase          =					torch.tensor(
				    [[-4_1.5_8_5_0, -4.4_4_4_0, -8.9_9_2_2], [-4_7.4_3_2_2, -5.9_1_4_3, -1.7_3_4_0]]   ,	device=lowercase			)
				assert torch.allclose(logits[0, :3, :3]   ,	lowercase   ,	atol=1e-4			)
		elif model_name == "blip2-flan-t5-xl-coco":
				UpperCamelCase          =					torch.tensor(
				    [[-5_7.0_1_0_9, -9.8_9_6_7, -1_2.6_2_8_0], [-6_8.6_5_7_8, -1_2.7_1_9_1, -1_0.5_0_6_5]]   ,	device=lowercase			)
		else:
				# cast to same type
				UpperCamelCase          =					logits.dtype
				assert torch.allclose(original_logits.to(lowercase			)   ,	lowercase   ,	atol=1e-2			)
		print('Looks ok!'			)

		print('Generating a caption...'			)
		UpperCamelCase          =					''
		UpperCamelCase          =					tokenizer(lowercase   ,	return_tensors='pt'			).input_ids.to(lowercase			)

		UpperCamelCase          =					original_model.generate({'image': original_pixel_values}			)
		UpperCamelCase          =					hf_model.generate(
		    lowercase   ,	lowercase   ,	do_sample=lowercase   ,	num_beams=5   ,	max_length=30   ,	min_length=1   ,	top_p=0.9   ,	repetition_penalty=1.0   ,	length_penalty=1.0   ,	temperature=1   ,	)
		print('Original generation:'   ,	lowercase			)
		UpperCamelCase          =					input_ids.shape[1]
		UpperCamelCase          =					processor.batch_decode(outputs[:, prompt_length:]   ,	skip_special_tokens=lowercase			)
		UpperCamelCase          =					[text.strip() for text in output_text]
		print('HF generation:'   ,	lowercase			)

		if pytorch_dump_folder_path is not None:
				processor.save_pretrained(lowercase			)
				hf_model.save_pretrained(lowercase			)

		if push_to_hub:
				processor.push_to_hub(f'''nielsr/{model_name}'''			)
				hf_model.push_to_hub(f'''nielsr/{model_name}'''			)


if __name__ == "__main__":
					_UpperCAmelCase :			Optional[int]              =	argparse.ArgumentParser()
					_UpperCAmelCase :			str              =	[
					    "blip2-opt-2.7b",
					    "blip2-opt-6.7b",
					    "blip2-opt-2.7b-coco",
					    "blip2-opt-6.7b-coco",
					    "blip2-flan-t5-xl",
					    "blip2-flan-t5-xl-coco",
					    "blip2-flan-t5-xxl",
					]
					parser.add_argument(
					    "--model_name",
					    default="blip2-opt-2.7b",
					    choices=choices,
					    type=str,
					    help="Path to hf config.json of model to convert",
					)
					parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
					parser.add_argument(
					    "--push_to_hub",
					    action="store_true",
					    help="Whether to push the model and processor to the hub after converting",
					)

					_UpperCAmelCase :			List[str]              =	parser.parse_args()

					convert_blipa_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
