'''simple docstring'''
import argparse
import datetime
import json
import time
import warnings
from logging import getLogger
from pathlib import Path
from typing import Dict, List

import torch
from tqdm import tqdm

from transformers import AutoModelForSeqaSeqLM, AutoTokenizer
from utils import calculate_bleu, calculate_rouge, chunks, parse_numeric_n_bool_cl_kwargs, use_task_specific_params


a=  getLogger(__name__)


a=  '''cuda''' if torch.cuda.is_available() else '''cpu'''







def  _UpperCamelCase (     _a	:	List[str]				,     _a	:	str				,     _a	:	str				,     _a	:	int = 8				,     _a	:	str = DEFAULT_DEVICE				,     _a	:	Union[str, Any]=False				,     _a	:	int="summarization"				,     _a	:	Optional[int]=None				,     **_a	:	Dict				,     ):




    """simple docstring"""



    __UpperCamelCase	:	Optional[Any]		    =  Path(_a			).open('w'				,     encoding='utf-8'			)
    __UpperCamelCase	:	Tuple		    =  str(_a			)
    __UpperCamelCase	:	List[str]		    =  AutoModelForSeqaSeqLM.from_pretrained(_a			).to(_a			)
    if fpaa:
        __UpperCamelCase	:	List[Any]		    =  model.half()

    __UpperCamelCase	:	Optional[int]		    =  AutoTokenizer.from_pretrained(_a			)
    logger.info(f"""Inferred tokenizer type: {tokenizer.__class__}"""			)  # if this is wrong, check config.model_type.

    __UpperCamelCase	:	Tuple		    =  time.time()
    # update config with task specific params
    use_task_specific_params(_a				,     _a			)
    if prefix is None:
        __UpperCamelCase	:	Tuple		    =  prefix or getattr(model.config				,     'prefix'				,     ''			) or ''
    for examples_chunk in tqdm(list(chunks(_a				,     _a			)			)			):
        __UpperCamelCase	:	str		    =  [prefix + text for text in examples_chunk]
        __UpperCamelCase	:	str		    =  tokenizer(_a				,     return_tensors='pt'				,     truncation=_a				,     padding='longest'			).to(_a			)
        __UpperCamelCase	:	Optional[Any]		    =  model.generate(
            input_ids=batch.input_ids				,     attention_mask=batch.attention_mask				,     **_a				,     )
        __UpperCamelCase	:	Tuple		    =  tokenizer.batch_decode(_a				,     skip_special_tokens=_a				,     clean_up_tokenization_spaces=_a			)
        for hypothesis in dec:
            fout.write(hypothesis + '\n'			)
            fout.flush()
    fout.close()
    __UpperCamelCase	:	Tuple		    =  int(time.time() - start_time			)  # seconds
    __UpperCamelCase	:	List[Any]		    =  len(_a			)
    return {"n_obs": n_obs, "runtime": runtime, "seconds_per_sample": round(runtime / n_obs				,     4			)}







def  _UpperCamelCase (     ):




    """simple docstring"""



    return datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'			)







def  _UpperCamelCase (     _a	:	List[Any]=True			):




    """simple docstring"""



    __UpperCamelCase	:	Union[str, Any]		    =  argparse.ArgumentParser()
    parser.add_argument('model_name'				,     type=_a				,     help='like facebook/bart-large-cnn,t5-base, etc.'			)
    parser.add_argument('input_path'				,     type=_a				,     help='like cnn_dm/test.source'			)
    parser.add_argument('save_path'				,     type=_a				,     help='where to save summaries'			)
    parser.add_argument('--reference_path'				,     type=_a				,     required=_a				,     help='like cnn_dm/test.target'			)
    parser.add_argument('--score_path'				,     type=_a				,     required=_a				,     default='metrics.json'				,     help='where to save metrics'			)
    parser.add_argument('--device'				,     type=_a				,     required=_a				,     default=_a				,     help='cuda, cuda:1, cpu etc.'			)
    parser.add_argument(
        '--prefix'				,     type=_a				,     required=_a				,     default=_a				,     help='will be added to the begininng of src examples'			)
    parser.add_argument('--task'				,     type=_a				,     default='summarization'				,     help='used for task_specific_params + metrics'			)
    parser.add_argument('--bs'				,     type=_a				,     default=8				,     required=_a				,     help='batch size'			)
    parser.add_argument(
        '--n_obs'				,     type=_a				,     default=-1				,     required=_a				,     help='How many observations. Defaults to all.'			)
    parser.add_argument('--fp16'				,     action='store_true'			)
    parser.add_argument('--dump-args'				,     action='store_true'				,     help='print the custom hparams with the results'			)
    parser.add_argument(
        '--info'				,     nargs='?'				,     type=_a				,     const=datetime_now()				,     help=(
            'use in conjunction w/ --dump-args to print with the results whatever other info you\'d like, e.g.'
            ' lang=en-ru. If no value is passed, the current datetime string will be used.'
        )				,     )
    # Unspecified args like --num_beams=2 --decoder_start_token_id=4 are passed to model.generate
    __UpperCamelCase	:	str		    =  parser.parse_known_args()
    __UpperCamelCase	:	str		    =  parse_numeric_n_bool_cl_kwargs(_a			)
    if parsed_args and verbose:
        print(f"""parsed the following generate kwargs: {parsed_args}"""			)
    __UpperCamelCase	:	List[Any]		    =  [' ' + x.rstrip() if 't5' in args.model_name else x.rstrip() for x in open(args.input_path			).readlines()]
    if args.n_obs > 0:
        __UpperCamelCase	:	int		    =  examples[: args.n_obs]
    Path(args.save_path			).parent.mkdir(exist_ok=_a			)

    if args.reference_path is None and Path(args.score_path			).exists():
        warnings.warn(f"""score_path {args.score_path} will be overwritten unless you type ctrl-c."""			)

    if args.device == "cpu" and args.fpaa:
        # this mix leads to RuntimeError: "threshold_cpu" not implemented for 'Half'
        raise ValueError('Can\'t mix --fp16 and --device cpu'			)

    __UpperCamelCase	:	str		    =  generate_summaries_or_translations(
        _a				,     args.save_path				,     args.model_name				,     batch_size=args.bs				,     device=args.device				,     fpaa=args.fpaa				,     task=args.task				,     prefix=args.prefix				,     **_a				,     )

    if args.reference_path is None:
        return {}

    # Compute scores
    __UpperCamelCase	:	List[Any]		    =  calculate_bleu if 'translation' in args.task else calculate_rouge
    __UpperCamelCase	:	Union[str, Any]		    =  [x.rstrip() for x in open(args.save_path			).readlines()]
    __UpperCamelCase	:	Union[str, Any]		    =  [x.rstrip() for x in open(args.reference_path			).readlines()][: len(_a			)]
    __UpperCamelCase	:	dict		    =  score_fn(_a				,     _a			)
    scores.update(_a			)

    if args.dump_args:
        scores.update(_a			)
    if args.info:
        __UpperCamelCase	:	Dict		    =  args.info

    if verbose:
        print(_a			)

    if args.score_path is not None:
        json.dump(_a				,     open(args.score_path				,     'w'			)			)

    return scores


if __name__ == "__main__":
  # Usage for MT:
  # python run_eval.py MODEL_NAME $DATA_DIR/test.source $save_dir/test_translations.txt --reference_path $DATA_DIR/test.target --score_path $save_dir/test_bleu.json  --task translation $@
  run_generate(verbose=True)






'''simple docstring'''
import argparse
import os

from . import (
    ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    BART_PRETRAINED_MODEL_ARCHIVE_LIST,
    BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    CAMEMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    CTRL_PRETRAINED_CONFIG_ARCHIVE_MAP,
    DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
    DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
    DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
    ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP,
    FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP,
    LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
    LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP,
    T5_PRETRAINED_CONFIG_ARCHIVE_MAP,
    TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP,
    WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP,
    XLM_PRETRAINED_CONFIG_ARCHIVE_MAP,
    XLM_ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP,
    XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP,
    AlbertConfig,
    BartConfig,
    BertConfig,
    CamembertConfig,
    CTRLConfig,
    DistilBertConfig,
    DPRConfig,
    ElectraConfig,
    FlaubertConfig,
    GPTaConfig,
    LayoutLMConfig,
    LxmertConfig,
    OpenAIGPTConfig,
    RobertaConfig,
    TaConfig,
    TFAlbertForPreTraining,
    TFBartForConditionalGeneration,
    TFBartForSequenceClassification,
    TFBertForPreTraining,
    TFBertForQuestionAnswering,
    TFBertForSequenceClassification,
    TFCamembertForMaskedLM,
    TFCTRLLMHeadModel,
    TFDistilBertForMaskedLM,
    TFDistilBertForQuestionAnswering,
    TFDPRContextEncoder,
    TFDPRQuestionEncoder,
    TFDPRReader,
    TFElectraForPreTraining,
    TFFlaubertWithLMHeadModel,
    TFGPTaLMHeadModel,
    TFLayoutLMForMaskedLM,
    TFLxmertForPreTraining,
    TFLxmertVisualFeatureEncoder,
    TFOpenAIGPTLMHeadModel,
    TFRobertaForCausalLM,
    TFRobertaForMaskedLM,
    TFRobertaForSequenceClassification,
    TFTaForConditionalGeneration,
    TFTransfoXLLMHeadModel,
    TFWavaVecaModel,
    TFXLMRobertaForMaskedLM,
    TFXLMWithLMHeadModel,
    TFXLNetLMHeadModel,
    TransfoXLConfig,
    WavaVecaConfig,
    WavaVecaModel,
    XLMConfig,
    XLMRobertaConfig,
    XLNetConfig,
    is_torch_available,
    load_pytorch_checkpoint_in_tfa_model,
)
from .utils import CONFIG_NAME, WEIGHTS_NAME, cached_file, logging


if is_torch_available():
  import numpy as np
  import torch

  from . import (
      AlbertForPreTraining,
      BartForConditionalGeneration,
      BertForPreTraining,
      BertForQuestionAnswering,
      BertForSequenceClassification,
      CamembertForMaskedLM,
      CTRLLMHeadModel,
      DistilBertForMaskedLM,
      DistilBertForQuestionAnswering,
      DPRContextEncoder,
      DPRQuestionEncoder,
      DPRReader,
      ElectraForPreTraining,
      FlaubertWithLMHeadModel,
      GPTaLMHeadModel,
      LayoutLMForMaskedLM,
      LxmertForPreTraining,
      LxmertVisualFeatureEncoder,
      OpenAIGPTLMHeadModel,
      RobertaForMaskedLM,
      RobertaForSequenceClassification,
      TaForConditionalGeneration,
      TransfoXLLMHeadModel,
      XLMRobertaForMaskedLM,
      XLMWithLMHeadModel,
      XLNetLMHeadModel,
  )


logging.set_verbosity_info()

a=  {
    '''bart''': (
        BartConfig,
        TFBartForConditionalGeneration,
        TFBartForSequenceClassification,
        BartForConditionalGeneration,
        BART_PRETRAINED_MODEL_ARCHIVE_LIST,
    ),
    '''bert''': (
        BertConfig,
        TFBertForPreTraining,
        BertForPreTraining,
        BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''bert-large-uncased-whole-word-masking-finetuned-squad''': (
        BertConfig,
        TFBertForQuestionAnswering,
        BertForQuestionAnswering,
        BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''bert-large-cased-whole-word-masking-finetuned-squad''': (
        BertConfig,
        TFBertForQuestionAnswering,
        BertForQuestionAnswering,
        BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''bert-base-cased-finetuned-mrpc''': (
        BertConfig,
        TFBertForSequenceClassification,
        BertForSequenceClassification,
        BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''dpr''': (
        DPRConfig,
        TFDPRQuestionEncoder,
        TFDPRContextEncoder,
        TFDPRReader,
        DPRQuestionEncoder,
        DPRContextEncoder,
        DPRReader,
        DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
    ),
    '''gpt2''': (
        GPTaConfig,
        TFGPTaLMHeadModel,
        GPTaLMHeadModel,
        GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''xlnet''': (
        XLNetConfig,
        TFXLNetLMHeadModel,
        XLNetLMHeadModel,
        XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''xlm''': (
        XLMConfig,
        TFXLMWithLMHeadModel,
        XLMWithLMHeadModel,
        XLM_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''xlm-roberta''': (
        XLMRobertaConfig,
        TFXLMRobertaForMaskedLM,
        XLMRobertaForMaskedLM,
        XLM_ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''transfo-xl''': (
        TransfoXLConfig,
        TFTransfoXLLMHeadModel,
        TransfoXLLMHeadModel,
        TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''openai-gpt''': (
        OpenAIGPTConfig,
        TFOpenAIGPTLMHeadModel,
        OpenAIGPTLMHeadModel,
        OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''roberta''': (
        RobertaConfig,
        TFRobertaForCausalLM,
        TFRobertaForMaskedLM,
        RobertaForMaskedLM,
        ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''layoutlm''': (
        LayoutLMConfig,
        TFLayoutLMForMaskedLM,
        LayoutLMForMaskedLM,
        LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
    ),
    '''roberta-large-mnli''': (
        RobertaConfig,
        TFRobertaForSequenceClassification,
        RobertaForSequenceClassification,
        ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''camembert''': (
        CamembertConfig,
        TFCamembertForMaskedLM,
        CamembertForMaskedLM,
        CAMEMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''flaubert''': (
        FlaubertConfig,
        TFFlaubertWithLMHeadModel,
        FlaubertWithLMHeadModel,
        FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''distilbert''': (
        DistilBertConfig,
        TFDistilBertForMaskedLM,
        DistilBertForMaskedLM,
        DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''distilbert-base-distilled-squad''': (
        DistilBertConfig,
        TFDistilBertForQuestionAnswering,
        DistilBertForQuestionAnswering,
        DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''lxmert''': (
        LxmertConfig,
        TFLxmertForPreTraining,
        LxmertForPreTraining,
        LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''lxmert-visual-feature-encoder''': (
        LxmertConfig,
        TFLxmertVisualFeatureEncoder,
        LxmertVisualFeatureEncoder,
        LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''ctrl''': (
        CTRLConfig,
        TFCTRLLMHeadModel,
        CTRLLMHeadModel,
        CTRL_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''albert''': (
        AlbertConfig,
        TFAlbertForPreTraining,
        AlbertForPreTraining,
        ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''t5''': (
        TaConfig,
        TFTaForConditionalGeneration,
        TaForConditionalGeneration,
        T5_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''electra''': (
        ElectraConfig,
        TFElectraForPreTraining,
        ElectraForPreTraining,
        ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''wav2vec2''': (
        WavaVecaConfig,
        TFWavaVecaModel,
        WavaVecaModel,
        WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
}







def  _UpperCamelCase (     _a	:	int				,     _a	:	Union[str, Any]				,     _a	:	str				,     _a	:	List[str]				,     _a	:	Dict=False				,     _a	:	str=True			):




    """simple docstring"""



    if model_type not in MODEL_CLASSES:
        raise ValueError(f"""Unrecognized model type, should be one of {list(MODEL_CLASSES.keys()			)}."""			)

    __UpperCamelCase    ,					__UpperCamelCase    ,					__UpperCamelCase    ,					__UpperCamelCase	:	Optional[int]		    =  MODEL_CLASSES[model_type]

    # Initialise TF model
    if config_file in aws_config_map:
        __UpperCamelCase	:	List[Any]		    =  cached_file(_a				,     _a				,     force_download=not use_cached_models			)
    __UpperCamelCase	:	Union[str, Any]		    =  config_class.from_json_file(_a			)
    __UpperCamelCase	:	Union[str, Any]		    =  True
    __UpperCamelCase	:	List[str]		    =  True
    print(f"""Building TensorFlow model from configuration: {config}"""			)
    __UpperCamelCase	:	Optional[int]		    =  model_class(_a			)

    # Load weights from tf checkpoint
    if pytorch_checkpoint_path in aws_config_map.keys():
        __UpperCamelCase	:	Union[str, Any]		    =  cached_file(
            _a				,     _a				,     force_download=not use_cached_models			)
    # Load PyTorch checkpoint in tf2 model:
    __UpperCamelCase	:	Any		    =  load_pytorch_checkpoint_in_tfa_model(_a				,     _a			)

    if compare_with_pt_model:
        __UpperCamelCase	:	Dict		    =  tf_model(tf_model.dummy_inputs				,     training=_a			)  # build the network

        __UpperCamelCase	:	List[str]		    =  torch.load(_a				,     map_location='cpu'			)
        __UpperCamelCase	:	Union[str, Any]		    =  pt_model_class.from_pretrained(
            pretrained_model_name_or_path=_a				,     config=_a				,     state_dict=_a			)

        with torch.no_grad():
            __UpperCamelCase	:	Any		    =  pt_model(**pt_model.dummy_inputs			)

        __UpperCamelCase	:	Optional[Any]		    =  pto[0].numpy()
        __UpperCamelCase	:	int		    =  tfo[0].numpy()
        __UpperCamelCase	:	Union[str, Any]		    =  np.amax(np.abs(np_pt - np_tf			)			)
        print(f"""Max absolute difference between models outputs {diff}"""			)
        assert diff <= 2E-2, f"""Error, model absolute difference is >2e-2: {diff}"""

    # Save pytorch-model
    print(f"""Save TensorFlow model to {tf_dump_path}"""			)
    tf_model.save_weights(_a				,     save_format='h5'			)







def  _UpperCamelCase (     _a	:	str				,     _a	:	Optional[Any]				,     _a	:	Dict=None				,     _a	:	Tuple=None				,     _a	:	int=False				,     _a	:	Union[str, Any]=False				,     _a	:	int=False				,     _a	:	List[str]=False				,     ):




    """simple docstring"""



    if args_model_type is None:
        __UpperCamelCase	:	Dict		    =  list(MODEL_CLASSES.keys()			)
    else:
        __UpperCamelCase	:	Union[str, Any]		    =  [args_model_type]

    for j, model_type in enumerate(_a				,     start=1			):
        print('=' * 1_0_0			)
        print(f""" Converting model type {j}/{len(_a			)}: {model_type}"""			)
        print('=' * 1_0_0			)
        if model_type not in MODEL_CLASSES:
            raise ValueError(f"""Unrecognized model type {model_type}, should be one of {list(MODEL_CLASSES.keys()			)}."""			)

        __UpperCamelCase    ,					__UpperCamelCase    ,					__UpperCamelCase    ,					__UpperCamelCase    ,					__UpperCamelCase	:	str		    =  MODEL_CLASSES[model_type]

        if model_shortcut_names_or_path is None:
            __UpperCamelCase	:	List[str]		    =  list(aws_model_maps.keys()			)
        if config_shortcut_names_or_path is None:
            __UpperCamelCase	:	Optional[int]		    =  model_shortcut_names_or_path

        for i, (model_shortcut_name, config_shortcut_name) in enumerate(
            zip(_a				,     _a			)				,     start=1			):
            print('-' * 1_0_0			)
            if "-squad" in model_shortcut_name or "-mrpc" in model_shortcut_name or "-mnli" in model_shortcut_name:
                if not only_convert_finetuned_models:
                    print(f"""    Skipping finetuned checkpoint {model_shortcut_name}"""			)
                    continue
                __UpperCamelCase	:	Dict		    =  model_shortcut_name
            elif only_convert_finetuned_models:
                print(f"""    Skipping not finetuned checkpoint {model_shortcut_name}"""			)
                continue
            print(
                f"""    Converting checkpoint {i}/{len(_a			)}: {model_shortcut_name} - model_type {model_type}"""			)
            print('-' * 1_0_0			)

            if config_shortcut_name in aws_config_map:
                __UpperCamelCase	:	Tuple		    =  cached_file(_a				,     _a				,     force_download=not use_cached_models			)
            else:
                __UpperCamelCase	:	Dict		    =  config_shortcut_name

            if model_shortcut_name in aws_model_maps:
                __UpperCamelCase	:	Dict		    =  cached_file(_a				,     _a				,     force_download=not use_cached_models			)
            else:
                __UpperCamelCase	:	Tuple		    =  model_shortcut_name

            if os.path.isfile(_a			):
                __UpperCamelCase	:	int		    =  'converted_model'

            convert_pt_checkpoint_to_tf(
                model_type=_a				,     pytorch_checkpoint_path=_a				,     config_file=_a				,     tf_dump_path=os.path.join(_a				,     model_shortcut_name + '-tf_model.h5'			)				,     compare_with_pt_model=_a				,     )
            if remove_cached_files:
                os.remove(_a			)
                os.remove(_a			)


if __name__ == "__main__":
  a=  argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '''--tf_dump_path''', default=None, type=str, required=True, help='''Path to the output Tensorflow dump file.'''
  )
  parser.add_argument(
      '''--model_type''',
      default=None,
      type=str,
      help=(
          F"""Model type selected in the list of {list(MODEL_CLASSES.keys())}. If not given, will download and """
          '''convert all the models from AWS.'''
      ),
  )
  parser.add_argument(
      '''--pytorch_checkpoint_path''',
      default=None,
      type=str,
      help=(
          '''Path to the PyTorch checkpoint path or shortcut name to download from AWS. '''
          '''If not given, will download and convert all the checkpoints from AWS.'''
      ),
  )
  parser.add_argument(
      '''--config_file''',
      default=None,
      type=str,
      help=(
          '''The config json file corresponding to the pre-trained model. \n'''
          '''This specifies the model architecture. If not given and '''
          '''--pytorch_checkpoint_path is not given or is a shortcut name '''
          '''use the configuration associated to the shortcut name on the AWS'''
      ),
  )
  parser.add_argument(
      '''--compare_with_pt_model''', action='''store_true''', help='''Compare Tensorflow and PyTorch model predictions.'''
  )
  parser.add_argument(
      '''--use_cached_models''',
      action='''store_true''',
      help='''Use cached models if possible instead of updating to latest checkpoint versions.''',
  )
  parser.add_argument(
      '''--remove_cached_files''',
      action='''store_true''',
      help='''Remove pytorch models after conversion (save memory when converting in batches).''',
  )
  parser.add_argument('''--only_convert_finetuned_models''', action='''store_true''', help='''Only convert finetuned models.''')
  a=  parser.parse_args()

  # if args.pytorch_checkpoint_path is not None:
  #     convert_pt_checkpoint_to_tf(args.model_type.lower(),
  #                                 args.pytorch_checkpoint_path,
  #                                 args.config_file if args.config_file is not None else args.pytorch_checkpoint_path,
  #                                 args.tf_dump_path,
  #                                 compare_with_pt_model=args.compare_with_pt_model,
  #                                 use_cached_models=args.use_cached_models)
  # else:
  convert_all_pt_checkpoints_to_tf(
      args.model_type.lower() if args.model_type is not None else None,
      args.tf_dump_path,
      model_shortcut_names_or_path=[args.pytorch_checkpoint_path]
      if args.pytorch_checkpoint_path is not None
      else None,
      config_shortcut_names_or_path=[args.config_file] if args.config_file is not None else None,
      compare_with_pt_model=args.compare_with_pt_model,
      use_cached_models=args.use_cached_models,
      remove_cached_files=args.remove_cached_files,
      only_convert_finetuned_models=args.only_convert_finetuned_models,
  )



