






"""simple docstring"""






import argparse
import json
import pickle
from pathlib import Path

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

from transformers import MaskFormerConfig, MaskFormerForInstanceSegmentation, MaskFormerImageProcessor, SwinConfig
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase_   : int   			= logging.get_logger(__name__)



def 					_A			(__a							) ->       List[str]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  int    						=		SwinConfig.from_pretrained(
        '''microsoft/swin-tiny-patch4-window7-224'''    ,				out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4''']							)
    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		MaskFormerConfig(backbone_config=__a							)

    SCREAMING_SNAKE_CASE_      :  List[Any]    						=		'''huggingface/label-files'''
    if "ade20k-full" in model_name:
        # this should be ok
        SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		8_47
        SCREAMING_SNAKE_CASE_      :  int    						=		'''maskformer-ade20k-full-id2label.json'''
    elif "ade" in model_name:
        # this should be ok
        SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		1_50
        SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		'''ade20k-id2label.json'''
    elif "coco-stuff" in model_name:
        # this should be ok
        SCREAMING_SNAKE_CASE_      :  Any    						=		1_71
        SCREAMING_SNAKE_CASE_      :  List[str]    						=		'''maskformer-coco-stuff-id2label.json'''
    elif "coco" in model_name:
        # TODO
        SCREAMING_SNAKE_CASE_      :  List[Any]    						=		1_33
        SCREAMING_SNAKE_CASE_      :  List[Any]    						=		'''coco-panoptic-id2label.json'''
    elif "cityscapes" in model_name:
        # this should be ok
        SCREAMING_SNAKE_CASE_      :  Dict    						=		19
        SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		'''cityscapes-id2label.json'''
    elif "vistas" in model_name:
        # this should be ok
        SCREAMING_SNAKE_CASE_      :  int    						=		65
        SCREAMING_SNAKE_CASE_      :  List[Any]    						=		'''mapillary-vistas-id2label.json'''

    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		json.load(open(hf_hub_download(__a    ,				__a    ,				repo_type='''dataset'''							)    ,				'''r'''							)							)
    SCREAMING_SNAKE_CASE_      :  Any    						=		{int(__a							): v for k, v in idalabel.items()}

    return config



def 					_A			(__a							) ->       Tuple:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  Any    						=		[]
    # stem
    # fmt: off
    rename_keys.append(('''backbone.patch_embed.proj.weight''', '''model.pixel_level_module.encoder.model.embeddings.patch_embeddings.projection.weight''')							)
    rename_keys.append(('''backbone.patch_embed.proj.bias''', '''model.pixel_level_module.encoder.model.embeddings.patch_embeddings.projection.bias''')							)
    rename_keys.append(('''backbone.patch_embed.norm.weight''', '''model.pixel_level_module.encoder.model.embeddings.norm.weight''')							)
    rename_keys.append(('''backbone.patch_embed.norm.bias''', '''model.pixel_level_module.encoder.model.embeddings.norm.bias''')							)
    # stages
    for i in range(len(config.backbone_config.depths							)							):
        for j in range(config.backbone_config.depths[i]							):
            rename_keys.append((f'backbone.layers.{i}.blocks.{j}.norm1.weight', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_before.weight')							)
            rename_keys.append((f'backbone.layers.{i}.blocks.{j}.norm1.bias', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_before.bias')							)
            rename_keys.append((f'backbone.layers.{i}.blocks.{j}.attn.relative_position_bias_table', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table')							)
            rename_keys.append((f'backbone.layers.{i}.blocks.{j}.attn.relative_position_index', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index')							)
            rename_keys.append((f'backbone.layers.{i}.blocks.{j}.attn.proj.weight', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight')							)
            rename_keys.append((f'backbone.layers.{i}.blocks.{j}.attn.proj.bias', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias')							)
            rename_keys.append((f'backbone.layers.{i}.blocks.{j}.norm2.weight', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_after.weight')							)
            rename_keys.append((f'backbone.layers.{i}.blocks.{j}.norm2.bias', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.layernorm_after.bias')							)
            rename_keys.append((f'backbone.layers.{i}.blocks.{j}.mlp.fc1.weight', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight')							)
            rename_keys.append((f'backbone.layers.{i}.blocks.{j}.mlp.fc1.bias', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias')							)
            rename_keys.append((f'backbone.layers.{i}.blocks.{j}.mlp.fc2.weight', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.output.dense.weight')							)
            rename_keys.append((f'backbone.layers.{i}.blocks.{j}.mlp.fc2.bias', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.blocks.{j}.output.dense.bias')							)

        if i < 3:
            rename_keys.append((f'backbone.layers.{i}.downsample.reduction.weight', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.reduction.weight')							)
            rename_keys.append((f'backbone.layers.{i}.downsample.norm.weight', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.norm.weight')							)
            rename_keys.append((f'backbone.layers.{i}.downsample.norm.bias', f'model.pixel_level_module.encoder.model.encoder.layers.{i}.downsample.norm.bias')							)
        rename_keys.append((f'backbone.norm{i}.weight', f'model.pixel_level_module.encoder.hidden_states_norms.{i}.weight')							)
        rename_keys.append((f'backbone.norm{i}.bias', f'model.pixel_level_module.encoder.hidden_states_norms.{i}.bias')							)

    # FPN
    rename_keys.append(('''sem_seg_head.layer_4.weight''', '''model.pixel_level_module.decoder.fpn.stem.0.weight''')							)
    rename_keys.append(('''sem_seg_head.layer_4.norm.weight''', '''model.pixel_level_module.decoder.fpn.stem.1.weight''')							)
    rename_keys.append(('''sem_seg_head.layer_4.norm.bias''', '''model.pixel_level_module.decoder.fpn.stem.1.bias''')							)
    for source_index, target_index in zip(range(3    ,				0    ,				-1							)    ,				range(0    ,				3							)							):
        rename_keys.append((f'sem_seg_head.adapter_{source_index}.weight', f'model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.0.weight')							)
        rename_keys.append((f'sem_seg_head.adapter_{source_index}.norm.weight', f'model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.1.weight')							)
        rename_keys.append((f'sem_seg_head.adapter_{source_index}.norm.bias', f'model.pixel_level_module.decoder.fpn.layers.{target_index}.proj.1.bias')							)
        rename_keys.append((f'sem_seg_head.layer_{source_index}.weight', f'model.pixel_level_module.decoder.fpn.layers.{target_index}.block.0.weight')							)
        rename_keys.append((f'sem_seg_head.layer_{source_index}.norm.weight', f'model.pixel_level_module.decoder.fpn.layers.{target_index}.block.1.weight')							)
        rename_keys.append((f'sem_seg_head.layer_{source_index}.norm.bias', f'model.pixel_level_module.decoder.fpn.layers.{target_index}.block.1.bias')							)
    rename_keys.append(('''sem_seg_head.mask_features.weight''', '''model.pixel_level_module.decoder.mask_projection.weight''')							)
    rename_keys.append(('''sem_seg_head.mask_features.bias''', '''model.pixel_level_module.decoder.mask_projection.bias''')							)

    # Transformer decoder
    for idx in range(config.decoder_config.decoder_layers							):
        # self-attention out projection
        rename_keys.append((f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.out_proj.weight', f'model.transformer_module.decoder.layers.{idx}.self_attn.out_proj.weight')							)
        rename_keys.append((f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.out_proj.bias', f'model.transformer_module.decoder.layers.{idx}.self_attn.out_proj.bias')							)
        # cross-attention out projection
        rename_keys.append((f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.out_proj.weight', f'model.transformer_module.decoder.layers.{idx}.encoder_attn.out_proj.weight')							)
        rename_keys.append((f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.out_proj.bias', f'model.transformer_module.decoder.layers.{idx}.encoder_attn.out_proj.bias')							)
        # MLP 1
        rename_keys.append((f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear1.weight', f'model.transformer_module.decoder.layers.{idx}.fc1.weight')							)
        rename_keys.append((f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear1.bias', f'model.transformer_module.decoder.layers.{idx}.fc1.bias')							)
        # MLP 2
        rename_keys.append((f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear2.weight', f'model.transformer_module.decoder.layers.{idx}.fc2.weight')							)
        rename_keys.append((f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.linear2.bias', f'model.transformer_module.decoder.layers.{idx}.fc2.bias')							)
        # layernorm 1 (self-attention layernorm)
        rename_keys.append((f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm1.weight', f'model.transformer_module.decoder.layers.{idx}.self_attn_layer_norm.weight')							)
        rename_keys.append((f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm1.bias', f'model.transformer_module.decoder.layers.{idx}.self_attn_layer_norm.bias')							)
        # layernorm 2 (cross-attention layernorm)
        rename_keys.append((f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm2.weight', f'model.transformer_module.decoder.layers.{idx}.encoder_attn_layer_norm.weight')							)
        rename_keys.append((f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm2.bias', f'model.transformer_module.decoder.layers.{idx}.encoder_attn_layer_norm.bias')							)
        # layernorm 3 (final layernorm)
        rename_keys.append((f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm3.weight', f'model.transformer_module.decoder.layers.{idx}.final_layer_norm.weight')							)
        rename_keys.append((f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.norm3.bias', f'model.transformer_module.decoder.layers.{idx}.final_layer_norm.bias')							)

    rename_keys.append(('''sem_seg_head.predictor.transformer.decoder.norm.weight''', '''model.transformer_module.decoder.layernorm.weight''')							)
    rename_keys.append(('''sem_seg_head.predictor.transformer.decoder.norm.bias''', '''model.transformer_module.decoder.layernorm.bias''')							)

    # heads on top
    rename_keys.append(('''sem_seg_head.predictor.query_embed.weight''', '''model.transformer_module.queries_embedder.weight''')							)

    rename_keys.append(('''sem_seg_head.predictor.input_proj.weight''', '''model.transformer_module.input_projection.weight''')							)
    rename_keys.append(('''sem_seg_head.predictor.input_proj.bias''', '''model.transformer_module.input_projection.bias''')							)

    rename_keys.append(('''sem_seg_head.predictor.class_embed.weight''', '''class_predictor.weight''')							)
    rename_keys.append(('''sem_seg_head.predictor.class_embed.bias''', '''class_predictor.bias''')							)

    for i in range(3							):
        rename_keys.append((f'sem_seg_head.predictor.mask_embed.layers.{i}.weight', f'mask_embedder.{i}.0.weight')							)
        rename_keys.append((f'sem_seg_head.predictor.mask_embed.layers.{i}.bias', f'mask_embedder.{i}.0.bias')							)
    # fmt: on

    return rename_keys



def 					_A			(__a    ,				__a    ,				__a							) ->       Optional[Any]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  str    						=		dct.pop(__a							)
    SCREAMING_SNAKE_CASE_      :  List[Any]    						=		val



def 					_A			(__a    ,				__a							) ->       Union[str, Any]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  Dict    						=		[int(backbone_config.embed_dim * 2**i							) for i in range(len(backbone_config.depths							)							)]
    for i in range(len(backbone_config.depths							)							):
        SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		num_features[i]
        for j in range(backbone_config.depths[i]							):
            # fmt: off
            # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
            SCREAMING_SNAKE_CASE_      :  Dict    						=		state_dict.pop(f'backbone.layers.{i}.blocks.{j}.attn.qkv.weight'							)
            SCREAMING_SNAKE_CASE_      :  List[Any]    						=		state_dict.pop(f'backbone.layers.{i}.blocks.{j}.attn.qkv.bias'							)
            # next, add query, keys and values (in that order) to the state dict
            SCREAMING_SNAKE_CASE_      :  List[str]    						=		in_proj_weight[:dim, :]
            SCREAMING_SNAKE_CASE_      :  int    						=		in_proj_bias[: dim]
            SCREAMING_SNAKE_CASE_      :  int    						=		in_proj_weight[
                dim : dim * 2, :
            ]
            SCREAMING_SNAKE_CASE_      :  Any    						=		in_proj_bias[
                dim : dim * 2
            ]
            SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		in_proj_weight[
                -dim :, :
            ]
            SCREAMING_SNAKE_CASE_      :  str    						=		in_proj_bias[-dim :]
            # fmt: on



def 					_A			(__a    ,				__a							) ->       Optional[Any]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  int    						=		config.decoder_config.hidden_size
    for idx in range(config.decoder_config.decoder_layers							):
        # read in weights + bias of self-attention input projection layer (in the original implementation, this is a single matrix + bias)
        SCREAMING_SNAKE_CASE_      :  List[Any]    						=		state_dict.pop(f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.in_proj_weight'							)
        SCREAMING_SNAKE_CASE_      :  Dict    						=		state_dict.pop(f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.self_attn.in_proj_bias'							)
        # next, add query, keys and values (in that order) to the state dict
        SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		in_proj_weight[: hidden_size, :]
        SCREAMING_SNAKE_CASE_      :  List[Any]    						=		in_proj_bias[:config.hidden_size]
        SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		in_proj_weight[hidden_size : hidden_size * 2, :]
        SCREAMING_SNAKE_CASE_      :  Dict    						=		in_proj_bias[hidden_size : hidden_size * 2]
        SCREAMING_SNAKE_CASE_      :  Dict    						=		in_proj_weight[-hidden_size :, :]
        SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		in_proj_bias[-hidden_size :]
        # read in weights + bias of cross-attention input projection layer (in the original implementation, this is a single matrix + bias)
        SCREAMING_SNAKE_CASE_      :  str    						=		state_dict.pop(f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.in_proj_weight'							)
        SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		state_dict.pop(f'sem_seg_head.predictor.transformer.decoder.layers.{idx}.multihead_attn.in_proj_bias'							)
        # next, add query, keys and values (in that order) to the state dict
        SCREAMING_SNAKE_CASE_      :  Any    						=		in_proj_weight[: hidden_size, :]
        SCREAMING_SNAKE_CASE_      :  List[Any]    						=		in_proj_bias[:config.hidden_size]
        SCREAMING_SNAKE_CASE_      :  Any    						=		in_proj_weight[hidden_size : hidden_size * 2, :]
        SCREAMING_SNAKE_CASE_      :  List[Any]    						=		in_proj_bias[hidden_size : hidden_size * 2]
        SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		in_proj_weight[-hidden_size :, :]
        SCREAMING_SNAKE_CASE_      :  Any    						=		in_proj_bias[-hidden_size :]
    # fmt: on



def 					_A			() ->       torch.Tensor:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  Any    						=		'''http://images.cocodataset.org/val2017/000000039769.jpg'''
    SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		Image.open(requests.get(__a    ,				stream=__a							).raw							)
    return im



@torch.no_grad()
def 					_A			(__a    ,				__a    ,				__a    ,				__a = False							) ->       Union[str, Any]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  int    						=		get_maskformer_config(__a							)

    # load original state_dict
    with open(__a    ,				'''rb'''							) as f:
        SCREAMING_SNAKE_CASE_      :  Any    						=		pickle.load(__a							)
    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		data['''model''']

    # for name, param in state_dict.items():
    #     print(name, param.shape)

    # rename keys
    SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		create_rename_keys(__a							)
    for src, dest in rename_keys:
        rename_key(__a    ,				__a    ,				__a							)
    read_in_swin_q_k_v(__a    ,				config.backbone_config							)
    read_in_decoder_q_k_v(__a    ,				__a							)

    # update to torch tensors
    for key, value in state_dict.items():
        SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		torch.from_numpy(__a							)

    # load 🤗 model
    SCREAMING_SNAKE_CASE_      :  Dict    						=		MaskFormerForInstanceSegmentation(__a							)
    model.eval()

    for name, param in model.named_parameters():
        print(__a    ,				param.shape							)

    SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_      :  List[str]    						=		model.load_state_dict(__a    ,				strict=__a							)
    assert missing_keys == [
        "model.pixel_level_module.encoder.model.layernorm.weight",
        "model.pixel_level_module.encoder.model.layernorm.bias",
    ]
    assert len(__a							) == 0, f'Unexpected keys: {unexpected_keys}'

    # verify results
    SCREAMING_SNAKE_CASE_      :  List[Any]    						=		prepare_img()
    if "vistas" in model_name:
        SCREAMING_SNAKE_CASE_      :  Tuple    						=		65
    elif "cityscapes" in model_name:
        SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		6_55_35
    else:
        SCREAMING_SNAKE_CASE_      :  int    						=		2_55
    SCREAMING_SNAKE_CASE_      :  Any    						=		True if '''ade''' in model_name else False
    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		MaskFormerImageProcessor(ignore_index=__a    ,				reduce_labels=__a							)

    SCREAMING_SNAKE_CASE_      :  int    						=		image_processor(__a    ,				return_tensors='''pt'''							)

    SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		model(**__a							)

    print('''Logits:'''    ,				outputs.class_queries_logits[0, :3, :3]							)

    if model_name == "maskformer-swin-tiny-ade":
        SCREAMING_SNAKE_CASE_      :  int    						=		torch.tensor(
            [[3.63_53, -4.47_70, -2.60_65], [0.50_81, -4.23_94, -3.53_43], [2.19_09, -5.03_53, -1.93_23]]							)
    assert torch.allclose(outputs.class_queries_logits[0, :3, :3]    ,				__a    ,				atol=1e-4							)
    print('''Looks ok!'''							)

    if pytorch_dump_folder_path is not None:
        print(f'Saving model and image processor to {pytorch_dump_folder_path}'							)
        Path(__a							).mkdir(exist_ok=__a							)
        model.save_pretrained(__a							)
        image_processor.save_pretrained(__a							)

    if push_to_hub:
        print('''Pushing model and image processor to the hub...'''							)
        model.push_to_hub(f'nielsr/{model_name}'							)
        image_processor.push_to_hub(f'nielsr/{model_name}'							)


if __name__ == "__main__":
  UpperCAmelCase_   : Optional[int]   			= argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--model_name""",
      default="""maskformer-swin-tiny-ade""",
      type=str,
      help=("""Name of the MaskFormer model you'd like to convert""",),
  )
  parser.add_argument(
      """--checkpoint_path""",
      default="""/Users/nielsrogge/Documents/MaskFormer_checkpoints/MaskFormer-Swin-tiny-ADE20k/model.pkl""",
      type=str,
      help="""Path to the original state dict (.pth file).""",
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
  )
  parser.add_argument(
      """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
  )

  UpperCAmelCase_   : Dict   			= parser.parse_args()
  convert_maskformer_checkpoint(
      args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
  )










"""simple docstring"""






import argparse
import logging
import pickle
from collections import Counter


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

if __name__ == "__main__":
  UpperCAmelCase_   : List[str]   			= argparse.ArgumentParser(
      description="""Token Counts for smoothing the masking probabilities in MLM (cf XLM/word2vec)"""
  )
  parser.add_argument(
      """--data_file""", type=str, default="""data/dump.bert-base-uncased.pickle""", help="""The binarized dataset."""
  )
  parser.add_argument(
      """--token_counts_dump""", type=str, default="""data/token_counts.bert-base-uncased.pickle""", help="""The dump file."""
  )
  parser.add_argument("""--vocab_size""", default=30522, type=int)
  UpperCAmelCase_   : Optional[Any]   			= parser.parse_args()

  logger.info(f'''Loading data from {args.data_file}''')
  with open(args.data_file, """rb""") as fp:
    UpperCAmelCase_   : Union[str, Any]   			= pickle.load(fp)

  logger.info("""Counting occurrences for MLM.""")
  UpperCAmelCase_   : Any   			= Counter()
  for tk_ids in data:
    counter.update(tk_ids)
  UpperCAmelCase_   : List[Any]   			= [0] * args.vocab_size
  for k, v in counter.items():
    UpperCAmelCase_   : Dict   			= v

  logger.info(f'''Dump to {args.token_counts_dump}''')
  with open(args.token_counts_dump, """wb""") as handle:
    pickle.dump(counts, handle, protocol=pickle.HIGHEST_PROTOCOL)


