




import os
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers.models.realm.configuration_realm import RealmConfig
from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer

class   __a(					_a       ):






    """simple docstring"""





    def 							a__				(				self  )			->		int:
         UpperCAmelCase_      :							Optional[int]       =      tempfile.mkdtemp()
         UpperCAmelCase_      :							str       =      5

         # Realm tok
         UpperCAmelCase_      :							List[str]       =      [
             '''[UNK]''',
             '''[CLS]''',
             '''[SEP]''',
             '''[PAD]''',
             '''[MASK]''',
             '''test''',
             '''question''',
             '''this''',
             '''is''',
             '''the''',
             '''first''',
             '''second''',
             '''third''',
             '''fourth''',
             '''fifth''',
             '''record''',
             '''want''',
             '''##want''',
             '''##ed''',
             '''wa''',
             '''un''',
             '''runn''',
             '''##ing''',
             ''',''',
             '''low''',
             '''lowest''',
         ]
         UpperCAmelCase_      :							int       =      os.path.join(self.tmpdirname       ,'''realm_tokenizer'''  )
         os.makedirs(_SCREAMING_SNAKE_CASE       ,exist_ok=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							str       =      os.path.join(_SCREAMING_SNAKE_CASE       ,VOCAB_FILES_NAMES['''vocab_file''']  )
         with open(self.vocab_file       ,'''w'''       ,encoding='''utf-8'''  ) as vocab_writer:
              vocab_writer.write(''''''.join([x + '''\n''' for x in vocab_tokens]  )  )

         UpperCAmelCase_      :							Optional[Any]       =      os.path.join(self.tmpdirname       ,'''realm_block_records'''  )
         os.makedirs(_SCREAMING_SNAKE_CASE       ,exist_ok=_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self  )			->		RealmTokenizer:
         return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname       ,'''realm_tokenizer'''  )  )





    def 							a__				(				self  )			->		Union[str, Any]:
         shutil.rmtree(self.tmpdirname  )





    def 							a__				(				self  )			->		int:
         UpperCAmelCase_      :							Dict       =      RealmConfig(num_block_records=self.num_block_records  )
         return config





    def 							a__				(				self  )			->		Tuple:
         UpperCAmelCase_      :							Any       =      Dataset.from_dict(
             {
                 '''id''': ['''0''', '''1'''],
                 '''question''': ['''foo''', '''bar'''],
                 '''answers''': [['''Foo''', '''Bar'''], ['''Bar''']],
             }  )
         return dataset





    def 							a__				(				self  )			->		Tuple:
         UpperCAmelCase_      :							str       =      np.array(
             [
                 b'''This is the first record''',
                 b'''This is the second record''',
                 b'''This is the third record''',
                 b'''This is the fourth record''',
                 b'''This is the fifth record''',
                 b'''This is a longer longer longer record''',
             ]       ,dtype=_SCREAMING_SNAKE_CASE       ,)
         return block_records





    def 							a__				(				self  )			->		List[str]:
         UpperCAmelCase_      :							Dict       =      RealmRetriever(
             block_records=self.get_dummy_block_records()       ,tokenizer=self.get_tokenizer()       ,)
         return retriever





    def 							a__				(				self  )			->		List[Any]:
         UpperCAmelCase_      :							Optional[Any]       =      self.get_config()
         UpperCAmelCase_      :							Optional[int]       =      self.get_dummy_retriever()
         UpperCAmelCase_      :							Optional[Any]       =      retriever.tokenizer

         UpperCAmelCase_      :							str       =      np.array([0, 3]       ,dtype='''long'''  )
         UpperCAmelCase_      :							Tuple       =      tokenizer(['''Test question''']  ).input_ids
         UpperCAmelCase_      :							Optional[Any]       =      tokenizer(
             ['''the fourth''']       ,add_special_tokens=_SCREAMING_SNAKE_CASE       ,return_token_type_ids=_SCREAMING_SNAKE_CASE       ,return_attention_mask=_SCREAMING_SNAKE_CASE       ,).input_ids
         UpperCAmelCase_      :							Any       =      config.reader_seq_len

         UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_      :							Optional[int]       =      retriever(
             _SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,answer_ids=_SCREAMING_SNAKE_CASE       ,max_length=_SCREAMING_SNAKE_CASE       ,return_tensors='''np'''  )

         self.assertEqual(len(_SCREAMING_SNAKE_CASE  )       ,2  )
         self.assertEqual(len(_SCREAMING_SNAKE_CASE  )       ,2  )
         self.assertEqual(len(_SCREAMING_SNAKE_CASE  )       ,2  )
         self.assertEqual(concat_inputs.input_ids.shape       ,(2, 10)  )
         self.assertEqual(concat_inputs.attention_mask.shape       ,(2, 10)  )
         self.assertEqual(concat_inputs.token_type_ids.shape       ,(2, 10)  )
         self.assertEqual(concat_inputs.special_tokens_mask.shape       ,(2, 10)  )
         self.assertEqual(
             tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0]  )       ,['''[CLS]''', '''test''', '''question''', '''[SEP]''', '''this''', '''is''', '''the''', '''first''', '''record''', '''[SEP]''']       ,)
         self.assertEqual(
             tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1]  )       ,['''[CLS]''', '''test''', '''question''', '''[SEP]''', '''this''', '''is''', '''the''', '''fourth''', '''record''', '''[SEP]''']       ,)





    def 							a__				(				self  )			->		str:
         UpperCAmelCase_      :							Union[str, Any]       =      self.get_config()
         UpperCAmelCase_      :							Optional[Any]       =      self.get_dummy_retriever()
         UpperCAmelCase_      :							List[Any]       =      retriever.tokenizer

         UpperCAmelCase_      :							Any       =      np.array([0, 3, 5]       ,dtype='''long'''  )
         UpperCAmelCase_      :							List[Any]       =      tokenizer(['''Test question''']  ).input_ids
         UpperCAmelCase_      :							Optional[Any]       =      tokenizer(
             ['''the fourth''', '''longer longer''']       ,add_special_tokens=_SCREAMING_SNAKE_CASE       ,return_token_type_ids=_SCREAMING_SNAKE_CASE       ,return_attention_mask=_SCREAMING_SNAKE_CASE       ,).input_ids
         UpperCAmelCase_      :							Any       =      config.reader_seq_len

         UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_      :							int       =      retriever(
             _SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,answer_ids=_SCREAMING_SNAKE_CASE       ,max_length=_SCREAMING_SNAKE_CASE       ,return_tensors='''np'''  )

         self.assertEqual([False, True, True]       ,_SCREAMING_SNAKE_CASE  )
         self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]       ,_SCREAMING_SNAKE_CASE  )
         self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]       ,_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self  )			->		Any:
         UpperCAmelCase_      :							List[Any]       =      self.get_dummy_retriever()
         retriever.save_pretrained(os.path.join(self.tmpdirname       ,'''realm_block_records'''  )  )

         # Test local path
         UpperCAmelCase_      :							Tuple       =      retriever.from_pretrained(os.path.join(self.tmpdirname       ,'''realm_block_records'''  )  )
         self.assertEqual(retriever.block_records[0]       ,b'''This is the first record'''  )

         # Test mocked remote path
         with patch('''transformers.models.realm.retrieval_realm.hf_hub_download'''  ) as mock_hf_hub_download:
              UpperCAmelCase_      :							Any       =      os.path.join(
                  os.path.join(self.tmpdirname       ,'''realm_block_records'''  )       ,_REALM_BLOCK_RECORDS_FILENAME  )
              UpperCAmelCase_      :							Union[str, Any]       =      RealmRetriever.from_pretrained('''google/realm-cc-news-pretrained-openqa'''  )

         self.assertEqual(retriever.block_records[0]       ,b'''This is the first record'''  )





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)