





"""simple docstring"""
import argparse
from collections import OrderedDict
from pathlib import Path

import torch

from transformers import (
    VisualBertConfig,
    VisualBertForMultipleChoice,
    VisualBertForPreTraining,
    VisualBertForQuestionAnswering,
    VisualBertForVisualReasoning,
)
from transformers.utils import logging


logging.set_verbosity_info()
__UpperCamelCase     :						Any						       =  logging.get_logger(__name__)

__UpperCamelCase     :						List[Any]						       =  [
    ('''bert.bert''', '''visual_bert'''),
    ('''bert.cls''', '''cls'''),
    ('''bert.classifier''', '''cls'''),
    ('''token_type_embeddings_visual''', '''visual_token_type_embeddings'''),
    ('''position_embeddings_visual''', '''visual_position_embeddings'''),
    ('''projection''', '''visual_projection'''),
]

__UpperCamelCase     :						List[str]						       =  [
    '''nlvr2_coco_pre_trained.th''',
    '''nlvr2_fine_tuned.th''',
    '''nlvr2_pre_trained.th''',
    '''vcr_coco_pre_train.th''',
    '''vcr_fine_tune.th''',
    '''vcr_pre_train.th''',
    '''vqa_coco_pre_trained.th''',
    '''vqa_fine_tuned.th''',
    '''vqa_pre_trained.th''',
]

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Union[str, Any]			):
 lowerCAmelCase    				= torch.load(_UpperCAmelCase  ,     map_location='cpu'			)
 return sd

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    List[str]  ,     _UpperCAmelCase					:    Union[str, Any]  ,     _UpperCAmelCase					:    Any=rename_keys_prefix			):
 lowerCAmelCase    				= OrderedDict()
 lowerCAmelCase    				= torch.arange(config.max_position_embeddings			).expand((1, -1)			)
 # detector_d = OrderedDict()
 for key in d:
  if "detector" in key:
   # detector_d[key.replace('detector.','')] = d[key]
   continue
  lowerCAmelCase    				= key
  for name_pair in rename_keys_prefix:
   lowerCAmelCase    				= new_key.replace(name_pair[0]  ,     name_pair[1]			)
  lowerCAmelCase    				= d[key]
  if key == "bert.cls.predictions.decoder.weight":
   # Old bert code didn't have `decoder.bias`, but was added separately
   lowerCAmelCase    				= new_d['cls.predictions.bias']
 return new_d

@torch.no_grad()
def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Tuple  ,     _UpperCAmelCase					:    Optional[Any]			):
 assert (
     checkpoint_path.split('/'			)[-1] in ACCEPTABLE_CHECKPOINTS
 ), F'The checkpoint provided must be in {ACCEPTABLE_CHECKPOINTS}.'

 # Get Config
 if "pre" in checkpoint_path:
  lowerCAmelCase    				= 'pretraining'
  if "vcr" in checkpoint_path:
   lowerCAmelCase    				= {'visual_embedding_dim': 512}
  elif "vqa_advanced" in checkpoint_path:
   lowerCAmelCase    				= {'visual_embedding_dim': 2048}
  elif "vqa" in checkpoint_path:
   lowerCAmelCase    				= {'visual_embedding_dim': 2048}
  elif "nlvr" in checkpoint_path:
   lowerCAmelCase    				= {'visual_embedding_dim': 1024}
  else:
   raise NotImplementedError(F'No implementation found for `{checkpoint_path}`.'			)
 else:
  if "vcr" in checkpoint_path:
   lowerCAmelCase    				= {'visual_embedding_dim': 512}
   lowerCAmelCase    				= 'multichoice'
  elif "vqa_advanced" in checkpoint_path:
   lowerCAmelCase    				= {'visual_embedding_dim': 2048}
   lowerCAmelCase    				= 'vqa_advanced'
  elif "vqa" in checkpoint_path:
   lowerCAmelCase    				= {'visual_embedding_dim': 2048, 'num_labels': 3129}
   lowerCAmelCase    				= 'vqa'
  elif "nlvr" in checkpoint_path:
   lowerCAmelCase    				= {
       'visual_embedding_dim': 1024,
       'num_labels': 2,
   }
   lowerCAmelCase    				= 'nlvr'

 lowerCAmelCase    				= VisualBertConfig(**_UpperCAmelCase			)

 # Load State Dict
 lowerCAmelCase    				= load_state_dict(_UpperCAmelCase			)

 lowerCAmelCase    				= get_new_dict(_UpperCAmelCase  ,     _UpperCAmelCase			)

 if model_type == "pretraining":
  lowerCAmelCase    				= VisualBertForPreTraining(_UpperCAmelCase			)
 elif model_type == "vqa":
  lowerCAmelCase    				= VisualBertForQuestionAnswering(_UpperCAmelCase			)
 elif model_type == "nlvr":
  lowerCAmelCase    				= VisualBertForVisualReasoning(_UpperCAmelCase			)
 elif model_type == "multichoice":
  lowerCAmelCase    				= VisualBertForMultipleChoice(_UpperCAmelCase			)

 model.load_state_dict(_UpperCAmelCase			)
 # Save Checkpoints
 Path(_UpperCAmelCase			).mkdir(exist_ok=_UpperCAmelCase			)
 model.save_pretrained(_UpperCAmelCase			)


if __name__ == "__main__":
       __UpperCamelCase     :						Optional[Any]						       =  argparse.ArgumentParser()
       # Required parameters
       parser.add_argument('''orig_checkpoint_path''', type=str, help='''A path to .th on local filesystem.''')
       parser.add_argument('''pytorch_dump_folder_path''', type=str, help='''Path to the output PyTorch model.''')
       __UpperCamelCase     :						Optional[Any]						       =  parser.parse_args()
       convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)












"""simple docstring"""
import argparse
import logging
import pickle
import random
import time

import numpy as np

from transformers import BertTokenizer, GPTaTokenizer, RobertaTokenizer


logging.basicConfig(
    format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s''', datefmt='''%m/%d/%Y %H:%M:%S''', level=logging.INFO
)
__UpperCamelCase     :						str						       =  logging.getLogger(__name__)

def 		_SCREAMING_SNAKE_CASE			():
 lowerCAmelCase    				= argparse.ArgumentParser(
     description='Preprocess the data to avoid re-doing it several times by (tokenization + token_to_ids).'			)
 parser.add_argument('--file_path'  ,     type=_UpperCAmelCase  ,     default='data/dump.txt'  ,     help='The path to the data.'			)
 parser.add_argument('--tokenizer_type'  ,     type=_UpperCAmelCase  ,     default='bert'  ,     choices=['bert', 'roberta', 'gpt2']			)
 parser.add_argument('--tokenizer_name'  ,     type=_UpperCAmelCase  ,     default='bert-base-uncased'  ,     help='The tokenizer to use.'			)
 parser.add_argument('--dump_file'  ,     type=_UpperCAmelCase  ,     default='data/dump'  ,     help='The dump file prefix.'			)
 lowerCAmelCase    				= parser.parse_args()

 logger.info(F'Loading Tokenizer ({args.tokenizer_name})'			)
 if args.tokenizer_type == "bert":
  lowerCAmelCase    				= BertTokenizer.from_pretrained(args.tokenizer_name			)
  lowerCAmelCase    				= tokenizer.special_tokens_map['cls_token']  # `[CLS]`
  lowerCAmelCase    				= tokenizer.special_tokens_map['sep_token']  # `[SEP]`
 elif args.tokenizer_type == "roberta":
  lowerCAmelCase    				= RobertaTokenizer.from_pretrained(args.tokenizer_name			)
  lowerCAmelCase    				= tokenizer.special_tokens_map['cls_token']  # `<s>`
  lowerCAmelCase    				= tokenizer.special_tokens_map['sep_token']  # `</s>`
 elif args.tokenizer_type == "gpt2":
  lowerCAmelCase    				= GPTaTokenizer.from_pretrained(args.tokenizer_name			)
  lowerCAmelCase    				= tokenizer.special_tokens_map['bos_token']  # `<|endoftext|>`
  lowerCAmelCase    				= tokenizer.special_tokens_map['eos_token']  # `<|endoftext|>`

 logger.info(F'Loading text from {args.file_path}'			)
 with open(args.file_path  ,     'r'  ,     encoding='utf8'			) as fp:
  lowerCAmelCase    				= fp.readlines()

 logger.info('Start encoding'			)
 logger.info(F'{len(_UpperCAmelCase			)} examples to process.'			)

 lowerCAmelCase    				= []
 lowerCAmelCase    				= 0
 lowerCAmelCase    				= 1_0000
 lowerCAmelCase    				= time.time()
 for text in data:
  lowerCAmelCase    				= F'{bos} {text.strip()} {sep}'
  lowerCAmelCase    				= tokenizer.encode(_UpperCAmelCase  ,     add_special_tokens=_UpperCAmelCase			)
  rslt.append(_UpperCAmelCase			)

  iter += 1
  if iter % interval == 0:
   lowerCAmelCase    				= time.time()
   logger.info(F'{iter} examples processed. - {(end-start):.2f}s/{interval}expl'			)
   lowerCAmelCase    				= time.time()
 logger.info('Finished binarization'			)
 logger.info(F'{len(_UpperCAmelCase			)} examples processed.'			)

 lowerCAmelCase    				= F'{args.dump_file}.{args.tokenizer_name}.pickle'
 lowerCAmelCase    				= tokenizer.vocab_size
 if vocab_size < (1 << 16):
  lowerCAmelCase    				= [np.uintaa(_UpperCAmelCase			) for d in rslt]
 else:
  lowerCAmelCase    				= [np.intaa(_UpperCAmelCase			) for d in rslt]
 random.shuffle(rslt_			)
 logger.info(F'Dump to {dp_file}'			)
 with open(_UpperCAmelCase  ,     'wb'			) as handle:
  pickle.dump(rslt_  ,     _UpperCAmelCase  ,     protocol=pickle.HIGHEST_PROTOCOL			)


if __name__ == "__main__":
       main()





