





import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    VideoMAEConfig,
    VideoMAEForPreTraining,
    VideoMAEForVideoClassification,
    VideoMAEImageProcessor,
)

def   a_   (			lowerCAmelCase_				:      Any     ):
   __lowerCAmelCase							     =     VideoMAEConfig()

   set_architecture_configs(lowerCAmelCase_,     lowerCAmelCase_     )

   if "finetuned" not in model_name:
      __lowerCAmelCase							     =     False

   if "finetuned" in model_name:
      __lowerCAmelCase							     =     '''huggingface/label-files'''
      if "kinetics" in model_name:
         __lowerCAmelCase							     =     400
         __lowerCAmelCase							     =     '''kinetics400-id2label.json'''
      elif "ssv2" in model_name:
         __lowerCAmelCase							     =     174
         __lowerCAmelCase							     =     '''something-something-v2-id2label.json'''
      else:
         raise ValueError('Model name should either contain \'kinetics\' or \'ssv2\' in case it\'s fine-tuned.'     )
      __lowerCAmelCase							     =     json.load(open(hf_hub_download(lowerCAmelCase_,     lowerCAmelCase_,     repo_type='dataset'     ),     'r'     )     )
      __lowerCAmelCase							     =     {int(lowerCAmelCase_     ): v for k, v in idalabel.items()}
      __lowerCAmelCase							     =     idalabel
      __lowerCAmelCase							     =     {v: k for k, v in idalabel.items()}

   return config

def   a_   (			lowerCAmelCase_				:      Tuple,     lowerCAmelCase_				:      Any     ):



   if "small" in model_name:
      __lowerCAmelCase							     =     384
      __lowerCAmelCase							     =     1536
      __lowerCAmelCase							     =     12
      __lowerCAmelCase							     =     16
      __lowerCAmelCase							     =     12
      __lowerCAmelCase							     =     3
      __lowerCAmelCase							     =     192
      __lowerCAmelCase							     =     768
   elif "large" in model_name:
      __lowerCAmelCase							     =     1024
      __lowerCAmelCase							     =     4096
      __lowerCAmelCase							     =     24
      __lowerCAmelCase							     =     16
      __lowerCAmelCase							     =     12
      __lowerCAmelCase							     =     8
      __lowerCAmelCase							     =     512
      __lowerCAmelCase							     =     2048
   elif "huge" in model_name:
      __lowerCAmelCase							     =     1280
      __lowerCAmelCase							     =     5120
      __lowerCAmelCase							     =     32
      __lowerCAmelCase							     =     16
      __lowerCAmelCase							     =     12
      __lowerCAmelCase							     =     8
      __lowerCAmelCase							     =     640
      __lowerCAmelCase							     =     2560
   elif "base" not in model_name:
      raise ValueError('Model name should include either "small", "base", "large", or "huge"'     )

def   a_   (			lowerCAmelCase_				:      Tuple     ):



   if "encoder." in name:
      __lowerCAmelCase							     =     name.replace('encoder.',     ''     )
   if "cls_token" in name:
      __lowerCAmelCase							     =     name.replace('cls_token',     'videomae.embeddings.cls_token'     )
   if "decoder_pos_embed" in name:
      __lowerCAmelCase							     =     name.replace('decoder_pos_embed',     'decoder.decoder_pos_embed'     )
   if "pos_embed" in name and "decoder" not in name:
      __lowerCAmelCase							     =     name.replace('pos_embed',     'videomae.embeddings.position_embeddings'     )
   if "patch_embed.proj" in name:
      __lowerCAmelCase							     =     name.replace('patch_embed.proj',     'videomae.embeddings.patch_embeddings.projection'     )
   if "patch_embed.norm" in name:
      __lowerCAmelCase							     =     name.replace('patch_embed.norm',     'videomae.embeddings.norm'     )
   if "decoder.blocks" in name:
      __lowerCAmelCase							     =     name.replace('decoder.blocks',     'decoder.decoder_layers'     )
   if "blocks" in name:
      __lowerCAmelCase							     =     name.replace('blocks',     'videomae.encoder.layer'     )
   if "attn.proj" in name:
      __lowerCAmelCase							     =     name.replace('attn.proj',     'attention.output.dense'     )
   if "attn" in name and "bias" not in name:
      __lowerCAmelCase							     =     name.replace('attn',     'attention.self'     )
   if "attn" in name:
      __lowerCAmelCase							     =     name.replace('attn',     'attention.attention'     )
   if "norm1" in name:
      __lowerCAmelCase							     =     name.replace('norm1',     'layernorm_before'     )
   if "norm2" in name:
      __lowerCAmelCase							     =     name.replace('norm2',     'layernorm_after'     )
   if "mlp.fc1" in name:
      __lowerCAmelCase							     =     name.replace('mlp.fc1',     'intermediate.dense'     )
   if "mlp.fc2" in name:
      __lowerCAmelCase							     =     name.replace('mlp.fc2',     'output.dense'     )
   if "decoder_embed" in name:
      __lowerCAmelCase							     =     name.replace('decoder_embed',     'decoder.decoder_embed'     )
   if "decoder_norm" in name:
      __lowerCAmelCase							     =     name.replace('decoder_norm',     'decoder.decoder_norm'     )
   if "decoder_pred" in name:
      __lowerCAmelCase							     =     name.replace('decoder_pred',     'decoder.decoder_pred'     )
   if "norm.weight" in name and "decoder" not in name and "fc" not in name:
      __lowerCAmelCase							     =     name.replace('norm.weight',     'videomae.layernorm.weight'     )
   if "norm.bias" in name and "decoder" not in name and "fc" not in name:
      __lowerCAmelCase							     =     name.replace('norm.bias',     'videomae.layernorm.bias'     )
   if "head" in name and "decoder" not in name:
      __lowerCAmelCase							     =     name.replace('head',     'classifier'     )

   return name

def   a_   (			lowerCAmelCase_				:      Tuple,     lowerCAmelCase_				:      List[str]     ):



   for key in orig_state_dict.copy().keys():
      __lowerCAmelCase							     =     orig_state_dict.pop(lowerCAmelCase_     )

      if key.startswith('encoder.'     ):
         __lowerCAmelCase							     =     key.replace('encoder.',     ''     )

      if "qkv" in key:
         __lowerCAmelCase							     =     key.split('.'     )
         if key.startswith('decoder.blocks'     ):
            __lowerCAmelCase							     =     config.decoder_hidden_size
            __lowerCAmelCase							     =     int(key_split[2]     )
            __lowerCAmelCase							     =     '''decoder.decoder_layers.'''
            if "weight" in key:
               __lowerCAmelCase							     =     val[:dim, :]
               __lowerCAmelCase							     =     val[dim : dim * 2, :]
               __lowerCAmelCase							     =     val[-dim:, :]
         else:
            __lowerCAmelCase							     =     config.hidden_size
            __lowerCAmelCase							     =     int(key_split[1]     )
            __lowerCAmelCase							     =     '''videomae.encoder.layer.'''
            if "weight" in key:
               __lowerCAmelCase							     =     val[:dim, :]
               __lowerCAmelCase							     =     val[dim : dim * 2, :]
               __lowerCAmelCase							     =     val[-dim:, :]
      else:
         __lowerCAmelCase							     =     val

   return orig_state_dict

def   a_   (			):
   __lowerCAmelCase							     =     hf_hub_download(
       repo_id='hf-internal-testing/spaghetti-video',     filename='eating_spaghetti.npy',     repo_type='dataset'     )
   __lowerCAmelCase							     =     np.load(lowerCAmelCase_     )
   return list(lowerCAmelCase_     )

def   a_   (			lowerCAmelCase_				:      Dict,     lowerCAmelCase_				:      Optional[int],     lowerCAmelCase_				:      Dict,     lowerCAmelCase_				:      Any     ):
   __lowerCAmelCase							     =     get_videomae_config(lowerCAmelCase_     )

   if "finetuned" in model_name:
      __lowerCAmelCase							     =     VideoMAEForVideoClassification(lowerCAmelCase_     )
   else:
      __lowerCAmelCase							     =     VideoMAEForPreTraining(lowerCAmelCase_     )

   # download original checkpoint, hosted on Google Drive
   __lowerCAmelCase							     =     '''pytorch_model.bin'''
   gdown.cached_download(lowerCAmelCase_,     lowerCAmelCase_,     quiet=lowerCAmelCase_     )
   __lowerCAmelCase							     =     torch.load(lowerCAmelCase_,     map_location='cpu'     )
   if "model" in files:
      __lowerCAmelCase							     =     files['''model''']
   else:
      __lowerCAmelCase							     =     files['''module''']
   __lowerCAmelCase							     =     convert_state_dict(lowerCAmelCase_,     lowerCAmelCase_     )

   model.load_state_dict(lowerCAmelCase_     )
   model.eval()

   # verify model on basic input
   __lowerCAmelCase							     =     VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5],     image_std=[0.5, 0.5, 0.5]     )
   __lowerCAmelCase							     =     prepare_video()
   __lowerCAmelCase							     =     image_processor(lowerCAmelCase_,     return_tensors='pt'     )

   if "finetuned" not in model_name:
      __lowerCAmelCase							     =     hf_hub_download(repo_id='hf-internal-testing/bool-masked-pos',     filename='bool_masked_pos.pt'     )
      __lowerCAmelCase							     =     torch.load(lowerCAmelCase_     )

   __lowerCAmelCase							     =     model(**lowerCAmelCase_     )
   __lowerCAmelCase							     =     outputs.logits

   __lowerCAmelCase							     =     [
       '''videomae-small-finetuned-kinetics''',
       '''videomae-small-finetuned-ssv2''',
       # Kinetics-400 checkpoints (short = pretrained only for 800 epochs instead of 1600)
       '''videomae-base-short''',
       '''videomae-base-short-finetuned-kinetics''',
       '''videomae-base''',
       '''videomae-base-finetuned-kinetics''',
       '''videomae-large''',
       '''videomae-large-finetuned-kinetics''',
       '''videomae-huge-finetuned-kinetics''',
       # Something-Something-v2 checkpoints (short = pretrained only for 800 epochs instead of 2400)
       '''videomae-base-short-ssv2''',
       '''videomae-base-short-finetuned-ssv2''',
       '''videomae-base-ssv2''',
       '''videomae-base-finetuned-ssv2''',
   ]

   # NOTE: logits were tested with image_mean and image_std equal to [0.5, 0.5, 0.5] and [0.5, 0.5, 0.5]
   if model_name == "videomae-small-finetuned-kinetics":
      __lowerCAmelCase							     =     torch.Size([1, 400]     )
      __lowerCAmelCase							     =     torch.tensor([-0.9291, -0.4061, -0.9307]     )
   elif model_name == "videomae-small-finetuned-ssv2":
      __lowerCAmelCase							     =     torch.Size([1, 174]     )
      __lowerCAmelCase							     =     torch.tensor([0.2671, -0.4689, -0.8235]     )
   elif model_name == "videomae-base":
      __lowerCAmelCase							     =     torch.Size([1, 1408, 1536]     )
      __lowerCAmelCase							     =     torch.tensor([[0.7739, 0.7968, 0.7089], [0.6701, 0.7487, 0.6209], [0.4287, 0.5158, 0.4773]]     )
   elif model_name == "videomae-base-short":
      __lowerCAmelCase							     =     torch.Size([1, 1408, 1536]     )
      __lowerCAmelCase							     =     torch.tensor([[0.7994, 0.9612, 0.8508], [0.7401, 0.8958, 0.8302], [0.5862, 0.7468, 0.7325]]     )
      # we verified the loss both for normalized and unnormalized targets for this one
      __lowerCAmelCase							     =     torch.tensor([0.5142]     ) if config.norm_pix_loss else torch.tensor([0.6469]     )
   elif model_name == "videomae-large":
      __lowerCAmelCase							     =     torch.Size([1, 1408, 1536]     )
      __lowerCAmelCase							     =     torch.tensor([[0.7149, 0.7997, 0.6966], [0.6768, 0.7869, 0.6948], [0.5139, 0.6221, 0.5605]]     )
   elif model_name == "videomae-large-finetuned-kinetics":
      __lowerCAmelCase							     =     torch.Size([1, 400]     )
      __lowerCAmelCase							     =     torch.tensor([0.0771, 0.0011, -0.3625]     )
   elif model_name == "videomae-huge-finetuned-kinetics":
      __lowerCAmelCase							     =     torch.Size([1, 400]     )
      __lowerCAmelCase							     =     torch.tensor([0.2433, 0.1632, -0.4894]     )
   elif model_name == "videomae-base-short-finetuned-kinetics":
      __lowerCAmelCase							     =     torch.Size([1, 400]     )
      __lowerCAmelCase							     =     torch.tensor([0.6588, 0.0990, -0.2493]     )
   elif model_name == "videomae-base-finetuned-kinetics":
      __lowerCAmelCase							     =     torch.Size([1, 400]     )
      __lowerCAmelCase							     =     torch.tensor([0.3669, -0.0688, -0.2421]     )
   elif model_name == "videomae-base-short-ssv2":
      __lowerCAmelCase							     =     torch.Size([1, 1408, 1536]     )
      __lowerCAmelCase							     =     torch.tensor([[0.4712, 0.5296, 0.5786], [0.2278, 0.2729, 0.4026], [0.0352, 0.0730, 0.2506]]     )
   elif model_name == "videomae-base-short-finetuned-ssv2":
      __lowerCAmelCase							     =     torch.Size([1, 174]     )
      __lowerCAmelCase							     =     torch.tensor([-0.0537, -0.1539, -0.3266]     )
   elif model_name == "videomae-base-ssv2":
      __lowerCAmelCase							     =     torch.Size([1, 1408, 1536]     )
      __lowerCAmelCase							     =     torch.tensor([[0.8131, 0.8727, 0.8546], [0.7366, 0.9377, 0.8870], [0.5935, 0.8874, 0.8564]]     )
   elif model_name == "videomae-base-finetuned-ssv2":
      __lowerCAmelCase							     =     torch.Size([1, 174]     )
      __lowerCAmelCase							     =     torch.tensor([0.1961, -0.8337, -0.6389]     )
   else:
      raise ValueError(F"""Model name not supported. Should be one of {model_names}"""     )

   # verify logits
   assert logits.shape == expected_shape
   if "finetuned" in model_name:
      assert torch.allclose(logits[0, :3],     lowerCAmelCase_,     atol=1E-4     )
   else:
      print('Logits:',     logits[0, :3, :3]     )
      assert torch.allclose(logits[0, :3, :3],     lowerCAmelCase_,     atol=1E-4     )
   print('Logits ok!'     )

   # verify loss, if applicable
   if model_name == "videomae-base-short":
      __lowerCAmelCase							     =     outputs.loss
      assert torch.allclose(lowerCAmelCase_,     lowerCAmelCase_,     atol=1E-4     )
      print('Loss ok!'     )

   if pytorch_dump_folder_path is not None:
      print(F"""Saving model and image processor to {pytorch_dump_folder_path}"""     )
      image_processor.save_pretrained(lowerCAmelCase_     )
      model.save_pretrained(lowerCAmelCase_     )

   if push_to_hub:
      print('Pushing to the hub...'     )
      model.push_to_hub(lowerCAmelCase_,     organization='nielsr'     )


if __name__ == "__main__":
  _snake_case  :   List[Any]						   =  argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '--checkpoint_url',
      default='https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4',
      type=str,
      help=(
          'URL of the original PyTorch checkpoint (on Google Drive) you\'d like to convert. Should be a direct'
          ' download link.'
      ),
  )
  parser.add_argument(
      '--pytorch_dump_folder_path',
      default='/Users/nielsrogge/Documents/VideoMAE/Test',
      type=str,
      help='Path to the output PyTorch model directory.',
  )
  parser.add_argument('--model_name', default='videomae-base', type=str, help='Name of the model.')
  parser.add_argument(
      '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
  )

  _snake_case  :   Union[str, Any]						   =  parser.parse_args()
  convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)








import argparse
import json
from typing import List

from ltp import LTP

from transformers.models.bert.tokenization_bert import BertTokenizer

def   a_   (			lowerCAmelCase_				:      Optional[Any]     ):
   # This defines a "chinese character" as anything in the CJK Unicode block:
   #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
   #
   # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
   # despite its name. The modern Korean Hangul alphabet is a different block,
   # as is Japanese Hiragana and Katakana. Those alphabets are used to write
   # space-separated words, so they are not treated specially and handled
   # like the all of the other languages.
   if (
       (cp >= 0X4e00 and cp <= 0X9fff)
       or (cp >= 0X3400 and cp <= 0X4dbf)  #
       or (cp >= 0X20000 and cp <= 0X2a6df)  #
       or (cp >= 0X2a700 and cp <= 0X2b73f)  #
       or (cp >= 0X2b740 and cp <= 0X2b81f)  #
       or (cp >= 0X2b820 and cp <= 0X2ceaf)  #
       or (cp >= 0Xf900 and cp <= 0Xfaff)
       or (cp >= 0X2f800 and cp <= 0X2fa1f)  #
   ):  #
      return True

   return False

def   a_   (			lowerCAmelCase_				:      str     ):
   # word like '180' or '身高' or '神'
   for char in word:
      __lowerCAmelCase							     =     ord(lowerCAmelCase_     )
      if not _is_chinese_char(lowerCAmelCase_     ):
         return 0
   return 1

def   a_   (			lowerCAmelCase_				:      List[str]     ):
   __lowerCAmelCase							     =     set()

   for token in tokens:
      __lowerCAmelCase							     =     len(lowerCAmelCase_     ) > 1 and is_chinese(lowerCAmelCase_     )
      if chinese_word:
         word_set.add(lowerCAmelCase_     )
   __lowerCAmelCase							     =     list(lowerCAmelCase_     )
   return word_list

def   a_   (			lowerCAmelCase_				:      List[str],     lowerCAmelCase_				:      set()     ):
   if not chinese_word_set:
      return bert_tokens
   __lowerCAmelCase							     =     max([len(lowerCAmelCase_     ) for w in chinese_word_set]     )

   __lowerCAmelCase							     =     bert_tokens
   __lowerCAmelCase     ,     __lowerCAmelCase							     =     0, len(lowerCAmelCase_     )
   while start < end:
      __lowerCAmelCase							     =     True
      if is_chinese(bert_word[start]     ):
         __lowerCAmelCase							     =     min(end - start,     lowerCAmelCase_     )
         for i in range(lowerCAmelCase_,     1,     -1     ):
            __lowerCAmelCase							     =     ''.join(bert_word[start : start + i]     )
            if whole_word in chinese_word_set:
               for j in range(start + 1,     start + i     ):
                  __lowerCAmelCase							     =     '##' + bert_word[j]
               __lowerCAmelCase							     =     start + i
               __lowerCAmelCase							     =     False
               break
      if single_word:
         start += 1
   return bert_word

def   a_   (			lowerCAmelCase_				:      List[str],     lowerCAmelCase_				:      LTP,     lowerCAmelCase_				:      BertTokenizer     ):
   __lowerCAmelCase							     =     []

   for i in range(0,     len(lowerCAmelCase_     ),     100     ):
      __lowerCAmelCase							     =     ltp_tokenizer.pipeline(lines[i : i + 100],     tasks=['cws']     ).cws
      __lowerCAmelCase							     =     [get_chinese_word(lowerCAmelCase_     ) for r in res]
      ltp_res.extend(lowerCAmelCase_     )
   assert len(lowerCAmelCase_     ) == len(lowerCAmelCase_     )

   __lowerCAmelCase							     =     []
   for i in range(0,     len(lowerCAmelCase_     ),     100     ):
      __lowerCAmelCase							     =     bert_tokenizer(lines[i : i + 100],     add_special_tokens=lowerCAmelCase_,     truncation=lowerCAmelCase_,     max_length=512     )
      bert_res.extend(res['input_ids']     )
   assert len(lowerCAmelCase_     ) == len(lowerCAmelCase_     )

   __lowerCAmelCase							     =     []
   for input_ids, chinese_word in zip(lowerCAmelCase_,     lowerCAmelCase_     ):
      __lowerCAmelCase							     =     []
      for id in input_ids:
         __lowerCAmelCase							     =     bert_tokenizer._convert_id_to_token(lowerCAmelCase_     )
         input_tokens.append(lowerCAmelCase_     )
      __lowerCAmelCase							     =     add_sub_symbol(lowerCAmelCase_,     lowerCAmelCase_     )
      __lowerCAmelCase							     =     []
      # We only save pos of chinese subwords start with ##, which mean is part of a whole word.
      for i, token in enumerate(lowerCAmelCase_     ):
         if token[:2] == "##":
            __lowerCAmelCase							     =     token[2:]
            # save chinese tokens' pos
            if len(lowerCAmelCase_     ) == 1 and _is_chinese_char(ord(lowerCAmelCase_     )     ):
               ref_id.append(lowerCAmelCase_     )
      ref_ids.append(lowerCAmelCase_     )

   assert len(lowerCAmelCase_     ) == len(lowerCAmelCase_     )

   return ref_ids

def   a_   (			lowerCAmelCase_				:      int     ):
   # For Chinese (Ro)Bert, the best result is from : RoBERTa-wwm-ext (https://github.com/ymcui/Chinese-BERT-wwm)
   # If we want to fine-tune these model, we have to use same tokenizer : LTP (https://github.com/HIT-SCIR/ltp)
   with open(args.file_name,     'r',     encoding='utf-8'     ) as f:
      __lowerCAmelCase							     =     f.readlines()
   __lowerCAmelCase							     =     [line.strip() for line in data if len(lowerCAmelCase_     ) > 0 and not line.isspace()]  # avoid delimiter like '\u2029'
   __lowerCAmelCase							     =     LTP(args.ltp     )  # faster in GPU device
   __lowerCAmelCase							     =     BertTokenizer.from_pretrained(args.bert     )

   __lowerCAmelCase							     =     prepare_ref(lowerCAmelCase_,     lowerCAmelCase_,     lowerCAmelCase_     )

   with open(args.save_path,     'w',     encoding='utf-8'     ) as f:
      __lowerCAmelCase							     =     [json.dumps(lowerCAmelCase_     ) + '\n' for ref in ref_ids]
      f.writelines(lowerCAmelCase_     )


if __name__ == "__main__":
     _snake_case  :   int						   =  argparse.ArgumentParser(description='prepare_chinese_ref')
     parser.add_argument(
         '--file_name',
         required=False,
         type=str,
         default='./resources/chinese-demo.txt',
         help='file need process, same as training data in lm',
     )
     parser.add_argument(
         '--ltp',
         required=False,
         type=str,
         default='./resources/ltp',
         help='resources for LTP tokenizer, usually a path',
     )
     parser.add_argument(
         '--bert',
         required=False,
         type=str,
         default='./resources/robert',
         help='resources for Bert tokenizer',
     )
     parser.add_argument(
         '--save_path',
         required=False,
         type=str,
         default='./resources/ref.txt',
         help='path to save res',
     )

     _snake_case  :   List[str]						   =  parser.parse_args()
     main(args)

