

"""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
)
SCREAMING_SNAKE_CASE_     :			Tuple	  =			logging.getLogger(__name__)



def 	UpperCAmelCase__   (				)      ->      str:




   """simple docstring"""
   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=A__	,    default="data/dump.txt"	,    help="The path to the data."   )
   parser.add_argument("--tokenizer_type"	,    type=A__	,    default="bert"	,    choices=["bert", "roberta", "gpt2"]   )
   parser.add_argument("--tokenizer_name"	,    type=A__	,    default="bert-base-uncased"	,    help="The tokenizer to use."   )
   parser.add_argument("--dump_file"	,    type=A__	,    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(A__   )} 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(A__	,    add_special_tokens=A__   )
      rslt.append(A__   )

      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(A__   )} examples processed.'   )

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


if __name__ == "__main__":
     main()







"""simple docstring"""




import unittest

from transformers import AlbertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         MODEL_FOR_PRETRAINING_MAPPING,
         AlbertForMaskedLM,
         AlbertForMultipleChoice,
         AlbertForPreTraining,
         AlbertForQuestionAnswering,
         AlbertForSequenceClassification,
         AlbertForTokenClassification,
         AlbertModel,
     )
     from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST


class 		_A       :





      def __init__(  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__=13    ,						SCREAMING_SNAKE_CASE__=7    ,						SCREAMING_SNAKE_CASE__=True    ,						SCREAMING_SNAKE_CASE__=True    ,						SCREAMING_SNAKE_CASE__=True    ,						SCREAMING_SNAKE_CASE__=True    ,						SCREAMING_SNAKE_CASE__=99    ,						SCREAMING_SNAKE_CASE__=16    ,						SCREAMING_SNAKE_CASE__=36    ,						SCREAMING_SNAKE_CASE__=6    ,						SCREAMING_SNAKE_CASE__=6    ,						SCREAMING_SNAKE_CASE__=6    ,						SCREAMING_SNAKE_CASE__=37    ,						SCREAMING_SNAKE_CASE__="gelu"    ,						SCREAMING_SNAKE_CASE__=0.1    ,						SCREAMING_SNAKE_CASE__=0.1    ,						SCREAMING_SNAKE_CASE__=512    ,						SCREAMING_SNAKE_CASE__=16    ,						SCREAMING_SNAKE_CASE__=2    ,						SCREAMING_SNAKE_CASE__=0.02    ,						SCREAMING_SNAKE_CASE__=3    ,						SCREAMING_SNAKE_CASE__=4    ,						SCREAMING_SNAKE_CASE__=None    ,						)   ->					Optional[Any]:
         lowerCamelCase__						=			parent
         lowerCamelCase__						=			batch_size
         lowerCamelCase__						=			seq_length
         lowerCamelCase__						=			is_training
         lowerCamelCase__						=			use_input_mask
         lowerCamelCase__						=			use_token_type_ids
         lowerCamelCase__						=			use_labels
         lowerCamelCase__						=			vocab_size
         lowerCamelCase__						=			embedding_size
         lowerCamelCase__						=			hidden_size
         lowerCamelCase__						=			num_hidden_layers
         lowerCamelCase__						=			num_hidden_groups
         lowerCamelCase__						=			num_attention_heads
         lowerCamelCase__						=			intermediate_size
         lowerCamelCase__						=			hidden_act
         lowerCamelCase__						=			hidden_dropout_prob
         lowerCamelCase__						=			attention_probs_dropout_prob
         lowerCamelCase__						=			max_position_embeddings
         lowerCamelCase__						=			type_vocab_size
         lowerCamelCase__						=			type_sequence_label_size
         lowerCamelCase__						=			initializer_range
         lowerCamelCase__						=			num_labels
         lowerCamelCase__						=			num_choices
         lowerCamelCase__						=			scope





      def 					_lowerCamelCase  (  self		)   ->					Any:
         lowerCamelCase__						=			ids_tensor([self.batch_size, self.seq_length]    ,						self.vocab_size		)

         lowerCamelCase__						=			None
         if self.use_input_mask:
            lowerCamelCase__						=			random_attention_mask([self.batch_size, self.seq_length]		)

         lowerCamelCase__						=			None
         if self.use_token_type_ids:
            lowerCamelCase__						=			ids_tensor([self.batch_size, self.seq_length]    ,						self.type_vocab_size		)

         lowerCamelCase__						=			None
         lowerCamelCase__						=			None
         lowerCamelCase__						=			None
         if self.use_labels:
            lowerCamelCase__						=			ids_tensor([self.batch_size]    ,						self.type_sequence_label_size		)
            lowerCamelCase__						=			ids_tensor([self.batch_size, self.seq_length]    ,						self.num_labels		)
            lowerCamelCase__						=			ids_tensor([self.batch_size]    ,						self.num_choices		)

         lowerCamelCase__						=			self.get_config()

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels





      def 					_lowerCamelCase  (  self		)   ->					List[str]:
         return AlbertConfig(
             vocab_size=self.vocab_size    ,						hidden_size=self.hidden_size    ,						num_hidden_layers=self.num_hidden_layers    ,						num_attention_heads=self.num_attention_heads    ,						intermediate_size=self.intermediate_size    ,						hidden_act=self.hidden_act    ,						hidden_dropout_prob=self.hidden_dropout_prob    ,						attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,						max_position_embeddings=self.max_position_embeddings    ,						type_vocab_size=self.type_vocab_size    ,						initializer_range=self.initializer_range    ,						num_hidden_groups=self.num_hidden_groups    ,						)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					int:
         lowerCamelCase__						=			AlbertModel(config=SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         lowerCamelCase__						=			model(SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			model(SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			model(SCREAMING_SNAKE_CASE__		)
         self.parent.assertEqual(result.last_hidden_state.shape    ,						(self.batch_size, self.seq_length, self.hidden_size)		)
         self.parent.assertEqual(result.pooler_output.shape    ,						(self.batch_size, self.hidden_size)		)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					str:
         lowerCamelCase__						=			AlbertForPreTraining(config=SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         lowerCamelCase__						=			model(
             SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__    ,						labels=SCREAMING_SNAKE_CASE__    ,						sentence_order_label=SCREAMING_SNAKE_CASE__    ,						)
         self.parent.assertEqual(result.prediction_logits.shape    ,						(self.batch_size, self.seq_length, self.vocab_size)		)
         self.parent.assertEqual(result.sop_logits.shape    ,						(self.batch_size, config.num_labels)		)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					Any:
         lowerCamelCase__						=			AlbertForMaskedLM(config=SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         lowerCamelCase__						=			model(SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__    ,						labels=SCREAMING_SNAKE_CASE__		)
         self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.vocab_size)		)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					Any:
         lowerCamelCase__						=			AlbertForQuestionAnswering(config=SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         lowerCamelCase__						=			model(
             SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__    ,						start_positions=SCREAMING_SNAKE_CASE__    ,						end_positions=SCREAMING_SNAKE_CASE__    ,						)
         self.parent.assertEqual(result.start_logits.shape    ,						(self.batch_size, self.seq_length)		)
         self.parent.assertEqual(result.end_logits.shape    ,						(self.batch_size, self.seq_length)		)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					Optional[int]:
         lowerCamelCase__						=			self.num_labels
         lowerCamelCase__						=			AlbertForSequenceClassification(SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         lowerCamelCase__						=			model(SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__    ,						labels=SCREAMING_SNAKE_CASE__		)
         self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_labels)		)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					int:
         lowerCamelCase__						=			self.num_labels
         lowerCamelCase__						=			AlbertForTokenClassification(config=SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         lowerCamelCase__						=			model(SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__    ,						labels=SCREAMING_SNAKE_CASE__		)
         self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.num_labels)		)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					int:
         lowerCamelCase__						=			self.num_choices
         lowerCamelCase__						=			AlbertForMultipleChoice(config=SCREAMING_SNAKE_CASE__		)
         model.to(SCREAMING_SNAKE_CASE__		)
         model.eval()
         lowerCamelCase__						=			input_ids.unsqueeze(1		).expand(-1    ,						self.num_choices    ,						-1		).contiguous()
         lowerCamelCase__						=			token_type_ids.unsqueeze(1		).expand(-1    ,						self.num_choices    ,						-1		).contiguous()
         lowerCamelCase__						=			input_mask.unsqueeze(1		).expand(-1    ,						self.num_choices    ,						-1		).contiguous()
         lowerCamelCase__						=			model(
             SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__    ,						token_type_ids=SCREAMING_SNAKE_CASE__    ,						labels=SCREAMING_SNAKE_CASE__    ,						)
         self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_choices)		)





      def 					_lowerCamelCase  (  self		)   ->					Union[str, Any]:
         lowerCamelCase__						=			self.prepare_config_and_inputs()
         (
             (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       
         )						=			config_and_inputs
         lowerCamelCase__						=			{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
         return config, inputs_dict






@require_torch
class 		_A       (			__a				,       __a				,       unittest.TestCase					):
      __a								=	(
          (
              AlbertModel,
              AlbertForPreTraining,
              AlbertForMaskedLM,
              AlbertForMultipleChoice,
              AlbertForSequenceClassification,
              AlbertForTokenClassification,
              AlbertForQuestionAnswering,
          )
          if is_torch_available()
          else ()
      )
      __a								=	(
          {
              'feature-extraction': AlbertModel,
              'fill-mask': AlbertForMaskedLM,
              'question-answering': AlbertForQuestionAnswering,
              'text-classification': AlbertForSequenceClassification,
              'token-classification': AlbertForTokenClassification,
              'zero-shot': AlbertForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      __a								=	True





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__=False		)   ->					Optional[Any]:
         lowerCamelCase__						=			super()._prepare_for_class(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						return_labels=SCREAMING_SNAKE_CASE__		)

         if return_labels:
            if model_class in get_values(SCREAMING_SNAKE_CASE__		):
               lowerCamelCase__						=			torch.zeros(
                   (self.model_tester.batch_size, self.model_tester.seq_length)    ,						dtype=torch.long    ,						device=SCREAMING_SNAKE_CASE__		)
               lowerCamelCase__						=			torch.zeros(
                   self.model_tester.batch_size    ,						dtype=torch.long    ,						device=SCREAMING_SNAKE_CASE__		)
         return inputs_dict





      def 					_lowerCamelCase  (  self		)   ->					Union[str, Any]:
         lowerCamelCase__						=			AlbertModelTester(self		)
         lowerCamelCase__						=			ConfigTester(self    ,						config_class=SCREAMING_SNAKE_CASE__    ,						hidden_size=37		)





      def 					_lowerCamelCase  (  self		)   ->					Any:
         self.config_tester.run_common_tests()





      def 					_lowerCamelCase  (  self		)   ->					int:
         lowerCamelCase__						=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self		)   ->					Optional[Any]:
         lowerCamelCase__						=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_pretraining(*SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self		)   ->					Dict:
         lowerCamelCase__						=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self		)   ->					Dict:
         lowerCamelCase__						=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self		)   ->					int:
         lowerCamelCase__						=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self		)   ->					Union[str, Any]:
         lowerCamelCase__						=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self		)   ->					Optional[int]:
         lowerCamelCase__						=			self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
            lowerCamelCase__						=			type
            self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__		)





      @slow
      def 					_lowerCamelCase  (  self		)   ->					int:
         for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            lowerCamelCase__						=			AlbertModel.from_pretrained(SCREAMING_SNAKE_CASE__		)
            self.assertIsNotNone(SCREAMING_SNAKE_CASE__		)






@require_torch
class 		_A       (			unittest.TestCase					):





      @slow
      def 					_lowerCamelCase  (  self		)   ->					int:
         lowerCamelCase__						=			AlbertModel.from_pretrained("albert-base-v2"		)
         lowerCamelCase__						=			torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]		)
         lowerCamelCase__						=			torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
         with torch.no_grad():
            lowerCamelCase__						=			model(SCREAMING_SNAKE_CASE__    ,						attention_mask=SCREAMING_SNAKE_CASE__		)[0]
         lowerCamelCase__						=			torch.Size((1, 11, 768)		)
         self.assertEqual(output.shape    ,						SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			torch.tensor(
             [[[-0.65_13, 1.50_35, -0.27_66], [-0.65_15, 1.50_46, -0.27_80], [-0.65_12, 1.50_49, -0.27_84]]]		)

         self.assertTrue(torch.allclose(output[:, 1:4, 1:4]    ,						SCREAMING_SNAKE_CASE__    ,						atol=1e-4		)		)




