




import json
import os
from functools import lru_cache
from typing import List, Optional, Tuple

import regex as re

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


__a          =			logging.get_logger(__name__)


__a          =			{'vocab_file': 'vocab.json', 'merges_file': 'merges.txt'}

__a          =			{
    'vocab_file': {
        'allenai/longformer-base-4096': 'https://huggingface.co/allenai/longformer-base-4096/resolve/main/vocab.json',
        'allenai/longformer-large-4096': (
            'https://huggingface.co/allenai/longformer-large-4096/resolve/main/vocab.json'
        ),
        'allenai/longformer-large-4096-finetuned-triviaqa': (
            'https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/vocab.json'
        ),
        'allenai/longformer-base-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/vocab.json'
        ),
        'allenai/longformer-large-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/vocab.json'
        ),
    },
    'merges_file': {
        'allenai/longformer-base-4096': 'https://huggingface.co/allenai/longformer-base-4096/resolve/main/merges.txt',
        'allenai/longformer-large-4096': (
            'https://huggingface.co/allenai/longformer-large-4096/resolve/main/merges.txt'
        ),
        'allenai/longformer-large-4096-finetuned-triviaqa': (
            'https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/merges.txt'
        ),
        'allenai/longformer-base-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/merges.txt'
        ),
        'allenai/longformer-large-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/merges.txt'
        ),
    },
}

__a          =			{
    'allenai/longformer-base-4096': 4_096,
    'allenai/longformer-large-4096': 4_096,
    'allenai/longformer-large-4096-finetuned-triviaqa': 4_096,
    'allenai/longformer-base-4096-extra.pos.embd.only': 4_096,
    'allenai/longformer-large-4096-extra.pos.embd.only': 4_096,
}







@lru_cache()
# Copied from transformers.models.roberta.tokenization_roberta.bytes_to_unicode
def lowerCamelCase__       (      ):

     '''simple docstring'''


     UpperCAmelCase_      :							int       =      (
         list(range(ord('''!'''				)			,						ord('''~'''				) + 1				)				) + list(range(ord('''¡'''				)			,						ord('''¬'''				) + 1				)				) + list(range(ord('''®'''				)			,						ord('''ÿ'''				) + 1				)				)
     )
     UpperCAmelCase_      :							Any       =      bs[:]
     UpperCAmelCase_      :							int       =      0
     for b in range(2**8				):
          if b not in bs:
               bs.append(_lowercase				)
               cs.append(2**8 + n				)
               n += 1
     UpperCAmelCase_      :							int       =      [chr(_lowercase				) for n in cs]
     return dict(zip(_lowercase			,						_lowercase				)				)







def lowerCamelCase__       (      _lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							Dict       =      set()
     UpperCAmelCase_      :							int       =      word[0]
     for char in word[1:]:
          pairs.add((prev_char, char)				)
          UpperCAmelCase_      :							Union[str, Any]       =      char
     return pairs



class   __a(					_a       ):






    """simple docstring"""





    lowerCAmelCase        =				VOCAB_FILES_NAMES
    lowerCAmelCase        =				PRETRAINED_VOCAB_FILES_MAP
    lowerCAmelCase        =				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    lowerCAmelCase        =				['''input_ids''', '''attention_mask''']





    def __init__(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE="replace"       ,_SCREAMING_SNAKE_CASE="<s>"       ,_SCREAMING_SNAKE_CASE="</s>"       ,_SCREAMING_SNAKE_CASE="</s>"       ,_SCREAMING_SNAKE_CASE="<s>"       ,_SCREAMING_SNAKE_CASE="<unk>"       ,_SCREAMING_SNAKE_CASE="<pad>"       ,_SCREAMING_SNAKE_CASE="<mask>"       ,_SCREAMING_SNAKE_CASE=False       ,**_SCREAMING_SNAKE_CASE       ,)			->		Tuple:
         UpperCAmelCase_      :							Union[str, Any]       =      AddedToken(_SCREAMING_SNAKE_CASE       ,lstrip=_SCREAMING_SNAKE_CASE       ,rstrip=_SCREAMING_SNAKE_CASE  ) if isinstance(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ) else bos_token
         UpperCAmelCase_      :							Any       =      AddedToken(_SCREAMING_SNAKE_CASE       ,lstrip=_SCREAMING_SNAKE_CASE       ,rstrip=_SCREAMING_SNAKE_CASE  ) if isinstance(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ) else eos_token
         UpperCAmelCase_      :							Optional[int]       =      AddedToken(_SCREAMING_SNAKE_CASE       ,lstrip=_SCREAMING_SNAKE_CASE       ,rstrip=_SCREAMING_SNAKE_CASE  ) if isinstance(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ) else sep_token
         UpperCAmelCase_      :							Optional[Any]       =      AddedToken(_SCREAMING_SNAKE_CASE       ,lstrip=_SCREAMING_SNAKE_CASE       ,rstrip=_SCREAMING_SNAKE_CASE  ) if isinstance(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ) else cls_token
         UpperCAmelCase_      :							Any       =      AddedToken(_SCREAMING_SNAKE_CASE       ,lstrip=_SCREAMING_SNAKE_CASE       ,rstrip=_SCREAMING_SNAKE_CASE  ) if isinstance(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ) else unk_token
         UpperCAmelCase_      :							int       =      AddedToken(_SCREAMING_SNAKE_CASE       ,lstrip=_SCREAMING_SNAKE_CASE       ,rstrip=_SCREAMING_SNAKE_CASE  ) if isinstance(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ) else pad_token

         # Mask token behave like a normal word, i.e. include the space before it
         UpperCAmelCase_      :							str       =      AddedToken(_SCREAMING_SNAKE_CASE       ,lstrip=_SCREAMING_SNAKE_CASE       ,rstrip=_SCREAMING_SNAKE_CASE  ) if isinstance(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ) else mask_token

         super().__init__(
             errors=_SCREAMING_SNAKE_CASE       ,bos_token=_SCREAMING_SNAKE_CASE       ,eos_token=_SCREAMING_SNAKE_CASE       ,unk_token=_SCREAMING_SNAKE_CASE       ,sep_token=_SCREAMING_SNAKE_CASE       ,cls_token=_SCREAMING_SNAKE_CASE       ,pad_token=_SCREAMING_SNAKE_CASE       ,mask_token=_SCREAMING_SNAKE_CASE       ,add_prefix_space=_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE       ,)

         with open(_SCREAMING_SNAKE_CASE       ,encoding='''utf-8'''  ) as vocab_handle:
              UpperCAmelCase_      :							Tuple       =      json.load(_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							int       =      {v: k for k, v in self.encoder.items()}
         UpperCAmelCase_      :							Optional[int]       =      errors  # how to handle errors in decoding
         UpperCAmelCase_      :							Optional[int]       =      bytes_to_unicode()
         UpperCAmelCase_      :							Dict       =      {v: k for k, v in self.byte_encoder.items()}
         with open(_SCREAMING_SNAKE_CASE       ,encoding='''utf-8'''  ) as merges_handle:
              UpperCAmelCase_      :							Any       =      merges_handle.read().split('''\n'''  )[1:-1]
         UpperCAmelCase_      :							Tuple       =      [tuple(merge.split()  ) for merge in bpe_merges]
         UpperCAmelCase_      :							Optional[Any]       =      dict(zip(_SCREAMING_SNAKE_CASE       ,range(len(_SCREAMING_SNAKE_CASE  )  )  )  )
         UpperCAmelCase_      :							List[str]       =      {}
         UpperCAmelCase_      :							str       =      add_prefix_space

         # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
         UpperCAmelCase_      :							Any       =      re.compile(R'''\'s|\'t|\'re|\'ve|\'m|\'ll|\'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+'''  )





    @property
    def 							a__				(				self  )			->		int:
         return len(self.encoder  )





    def 							a__				(				self  )			->		Optional[Any]:
         return dict(self.encoder       ,**self.added_tokens_encoder  )





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE  )			->		int:
         if token in self.cache:
              return self.cache[token]
         UpperCAmelCase_      :							int       =      tuple(_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Any       =      get_pairs(_SCREAMING_SNAKE_CASE  )

         if not pairs:
              return token

         while True:
              UpperCAmelCase_      :							Dict       =      min(_SCREAMING_SNAKE_CASE       ,key=lambda _SCREAMING_SNAKE_CASE  : self.bpe_ranks.get(_SCREAMING_SNAKE_CASE       ,float('''inf'''  )  )  )
              if bigram not in self.bpe_ranks:
                   break
              UpperCAmelCase_,   UpperCAmelCase_      :							Optional[int]       =      bigram
              UpperCAmelCase_      :							Optional[int]       =      []
              UpperCAmelCase_      :							Any       =      0
              while i < len(_SCREAMING_SNAKE_CASE  ):
                   try:
                        UpperCAmelCase_      :							Dict       =      word.index(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
                   except ValueError:
                        new_word.extend(word[i:]  )
                        break
                   else:
                        new_word.extend(word[i:j]  )
                        UpperCAmelCase_      :							Tuple       =      j

                   if word[i] == first and i < len(_SCREAMING_SNAKE_CASE  ) - 1 and word[i + 1] == second:
                        new_word.append(first + second  )
                        i += 2
                   else:
                        new_word.append(word[i]  )
                        i += 1
              UpperCAmelCase_      :							List[Any]       =      tuple(_SCREAMING_SNAKE_CASE  )
              UpperCAmelCase_      :							Union[str, Any]       =      new_word
              if len(_SCREAMING_SNAKE_CASE  ) == 1:
                   break
              else:
                   UpperCAmelCase_      :							Optional[int]       =      get_pairs(_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Dict       =      ''' '''.join(_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							List[Any]       =      word
         return word





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE  )			->		Union[str, Any]:
         UpperCAmelCase_      :							Dict       =      []
         for token in re.findall(self.pat       ,_SCREAMING_SNAKE_CASE  ):
              UpperCAmelCase_      :							Any       =      ''''''.join(
                  self.byte_encoder[b] for b in token.encode('''utf-8'''  )  )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
              bpe_tokens.extend(bpe_token for bpe_token in self.bpe(_SCREAMING_SNAKE_CASE  ).split(''' '''  )  )
         return bpe_tokens





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE  )			->		Tuple:
         return self.encoder.get(_SCREAMING_SNAKE_CASE       ,self.encoder.get(self.unk_token  )  )





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE  )			->		Tuple:
         return self.decoder.get(_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE  )			->		int:
         UpperCAmelCase_      :							Optional[Any]       =      ''''''.join(_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Union[str, Any]       =      bytearray([self.byte_decoder[c] for c in text]  ).decode('''utf-8'''       ,errors=self.errors  )
         return text





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE = None  )			->		Tuple[str]:
         if not os.path.isdir(_SCREAMING_SNAKE_CASE  ):
              logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''  )
              return
         UpperCAmelCase_      :							Dict       =      os.path.join(
             _SCREAMING_SNAKE_CASE       ,(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']  )
         UpperCAmelCase_      :							Tuple       =      os.path.join(
             _SCREAMING_SNAKE_CASE       ,(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''merges_file''']  )

         with open(_SCREAMING_SNAKE_CASE       ,'''w'''       ,encoding='''utf-8'''  ) as f:
              f.write(json.dumps(self.encoder       ,indent=2       ,sort_keys=_SCREAMING_SNAKE_CASE       ,ensure_ascii=_SCREAMING_SNAKE_CASE  ) + '''\n'''  )

         UpperCAmelCase_      :							Optional[Any]       =      0
         with open(_SCREAMING_SNAKE_CASE       ,'''w'''       ,encoding='''utf-8'''  ) as writer:
              writer.write('''#version: 0.2\n'''  )
              for bpe_tokens, token_index in sorted(self.bpe_ranks.items()       ,key=lambda _SCREAMING_SNAKE_CASE  : kv[1]  ):
                   if index != token_index:
                        logger.warning(
                            f'''Saving vocabulary to {merge_file}: BPE merge indices are not consecutive.'''
                            ''' Please check that the tokenizer is not corrupted!'''  )
                        UpperCAmelCase_      :							List[str]       =      token_index
                   writer.write(''' '''.join(_SCREAMING_SNAKE_CASE  ) + '''\n'''  )
                   index += 1

         return vocab_file, merge_file





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE = None  )			->		List[int]:
         if token_ids_a is None:
              return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
         UpperCAmelCase_      :							Dict       =      [self.cls_token_id]
         UpperCAmelCase_      :							Union[str, Any]       =      [self.sep_token_id]
         return cls + token_ids_a + sep + sep + token_ids_a + sep





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE = None       ,_SCREAMING_SNAKE_CASE = False  )			->		List[int]:
         if already_has_special_tokens:
              return super().get_special_tokens_mask(
                  token_ids_a=_SCREAMING_SNAKE_CASE       ,token_ids_a=_SCREAMING_SNAKE_CASE       ,already_has_special_tokens=_SCREAMING_SNAKE_CASE  )

         if token_ids_a is None:
              return [1] + ([0] * len(_SCREAMING_SNAKE_CASE  )) + [1]
         return [1] + ([0] * len(_SCREAMING_SNAKE_CASE  )) + [1, 1] + ([0] * len(_SCREAMING_SNAKE_CASE  )) + [1]





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE = None  )			->		List[int]:
         UpperCAmelCase_      :							Tuple       =      [self.sep_token_id]
         UpperCAmelCase_      :							List[Any]       =      [self.cls_token_id]

         if token_ids_a is None:
              return len(cls + token_ids_a + sep  ) * [0]
         return len(cls + token_ids_a + sep + sep + token_ids_a + sep  ) * [0]





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE=False       ,**_SCREAMING_SNAKE_CASE  )			->		List[Any]:
         UpperCAmelCase_      :							str       =      kwargs.pop('''add_prefix_space'''       ,self.add_prefix_space  )
         if (is_split_into_words or add_prefix_space) and (len(_SCREAMING_SNAKE_CASE  ) > 0 and not text[0].isspace()):
              UpperCAmelCase_      :							List[Any]       =      ''' ''' + text
         return (text, kwargs)





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
)
__a          =			logging.getLogger(__name__)

if __name__ == "__main__":
     __a          =			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=30_522, type=int)
     __a          =			parser.parse_args()

     logger.info(F"""Loading data from {args.data_file}""")
     with open(args.data_file, 'rb') as fp:
          __a          =			pickle.load(fp)

     logger.info('Counting occurrences for MLM.')
     __a          =			Counter()
     for tk_ids in data:
          counter.update(tk_ids)
     __a          =			[0] * args.vocab_size
     for k, v in counter.items():
          __a          =			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)