





'''simple docstring'''






import unittest

from transformers import SqueezeBertConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, 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 (
          SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          SqueezeBertForMaskedLM,
          SqueezeBertForMultipleChoice,
          SqueezeBertForQuestionAnswering,
          SqueezeBertForSequenceClassification,
          SqueezeBertForTokenClassification,
          SqueezeBertModel,
      )




class 					UpperCAmelCase	(     snake_case_  ):




      def __init__(     self     :					Union[str, Any]     ,     __snake_case     :					Any     ,     __snake_case     :					Any=13     ,     __snake_case     :					Optional[int]=7     ,     __snake_case     :					Optional[Any]=True     ,     __snake_case     :					str=True     ,     __snake_case     :					int=False     ,     __snake_case     :					int=True     ,     __snake_case     :					Any=99     ,     __snake_case     :					List[Any]=32     ,     __snake_case     :					Dict=5     ,     __snake_case     :					Union[str, Any]=4     ,     __snake_case     :					Optional[int]=64     ,     __snake_case     :					Any="gelu"     ,     __snake_case     :					Tuple=0.1     ,     __snake_case     :					str=0.1     ,     __snake_case     :					Optional[int]=5_12     ,     __snake_case     :					Dict=16     ,     __snake_case     :					Optional[int]=2     ,     __snake_case     :					List[Any]=0.02     ,     __snake_case     :					int=3     ,     __snake_case     :					Optional[int]=4     ,     __snake_case     :					str=None     ,     __snake_case     :					Any=2     ,     __snake_case     :					Union[str, Any]=2     ,     __snake_case     :					List[str]=2     ,     __snake_case     :					Tuple=2     ,     __snake_case     :					Dict=4     ,     __snake_case     :					Optional[Any]=1     ,     )							-> str:
            _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					=				hidden_size
            _lowerCAmelCase					=				num_hidden_layers
            _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
            _lowerCAmelCase					=				q_groups
            _lowerCAmelCase					=				k_groups
            _lowerCAmelCase					=				v_groups
            _lowerCAmelCase					=				post_attention_groups
            _lowerCAmelCase					=				intermediate_groups
            _lowerCAmelCase					=				output_groups




      def 							lowercase__       (     self     :					int    )							-> str:
            _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
            _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, input_mask, sequence_labels, token_labels, choice_labels




      def 							lowercase__       (     self     :					Any    )							-> Optional[Any]:
            return SqueezeBertConfig(
                embedding_size=self.hidden_size     ,     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     ,     attention_probs_dropout_prob=self.hidden_dropout_prob     ,     attention_dropout=self.attention_probs_dropout_prob     ,     max_position_embeddings=self.max_position_embeddings     ,     initializer_range=self.initializer_range     ,     q_groups=self.q_groups     ,     k_groups=self.k_groups     ,     v_groups=self.v_groups     ,     post_attention_groups=self.post_attention_groups     ,     intermediate_groups=self.intermediate_groups     ,     output_groups=self.output_groups     ,     )




      def 							lowercase__       (     self     :					Tuple     ,     __snake_case     :					Tuple     ,     __snake_case     :					Union[str, Any]     ,     __snake_case     :					Any     ,     __snake_case     :					List[Any]     ,     __snake_case     :					List[Any]     ,     __snake_case     :					Dict    )							-> Optional[int]:
            _lowerCAmelCase					=				SqueezeBertModel(config=__snake_case    )
            model.to(__snake_case    )
            model.eval()
            _lowerCAmelCase					=				model(__snake_case     ,     __snake_case    )
            _lowerCAmelCase					=				model(__snake_case    )
            self.parent.assertEqual(result.last_hidden_state.shape     ,     (self.batch_size, self.seq_length, self.hidden_size)    )




      def 							lowercase__       (     self     :					Union[str, Any]     ,     __snake_case     :					Optional[Any]     ,     __snake_case     :					Dict     ,     __snake_case     :					Union[str, Any]     ,     __snake_case     :					int     ,     __snake_case     :					Union[str, Any]     ,     __snake_case     :					List[Any]    )							-> List[Any]:
            _lowerCAmelCase					=				SqueezeBertForMaskedLM(config=__snake_case    )
            model.to(__snake_case    )
            model.eval()
            _lowerCAmelCase					=				model(__snake_case     ,     attention_mask=__snake_case     ,     labels=__snake_case    )
            self.parent.assertEqual(result.logits.shape     ,     (self.batch_size, self.seq_length, self.vocab_size)    )




      def 							lowercase__       (     self     :					Dict     ,     __snake_case     :					Tuple     ,     __snake_case     :					Any     ,     __snake_case     :					Dict     ,     __snake_case     :					Optional[Any]     ,     __snake_case     :					Optional[int]     ,     __snake_case     :					int    )							-> Optional[int]:
            _lowerCAmelCase					=				SqueezeBertForQuestionAnswering(config=__snake_case    )
            model.to(__snake_case    )
            model.eval()
            _lowerCAmelCase					=				model(
                __snake_case     ,     attention_mask=__snake_case     ,     start_positions=__snake_case     ,     end_positions=__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 							lowercase__       (     self     :					Tuple     ,     __snake_case     :					int     ,     __snake_case     :					Tuple     ,     __snake_case     :					Optional[int]     ,     __snake_case     :					Dict     ,     __snake_case     :					Tuple     ,     __snake_case     :					Union[str, Any]    )							-> Union[str, Any]:
            _lowerCAmelCase					=				self.num_labels
            _lowerCAmelCase					=				SqueezeBertForSequenceClassification(__snake_case    )
            model.to(__snake_case    )
            model.eval()
            _lowerCAmelCase					=				model(__snake_case     ,     attention_mask=__snake_case     ,     labels=__snake_case    )
            self.parent.assertEqual(result.logits.shape     ,     (self.batch_size, self.num_labels)    )




      def 							lowercase__       (     self     :					int     ,     __snake_case     :					Union[str, Any]     ,     __snake_case     :					Union[str, Any]     ,     __snake_case     :					Optional[Any]     ,     __snake_case     :					Dict     ,     __snake_case     :					Any     ,     __snake_case     :					Optional[int]    )							-> int:
            _lowerCAmelCase					=				self.num_labels
            _lowerCAmelCase					=				SqueezeBertForTokenClassification(config=__snake_case    )
            model.to(__snake_case    )
            model.eval()

            _lowerCAmelCase					=				model(__snake_case     ,     attention_mask=__snake_case     ,     labels=__snake_case    )
            self.parent.assertEqual(result.logits.shape     ,     (self.batch_size, self.seq_length, self.num_labels)    )




      def 							lowercase__       (     self     :					Optional[int]     ,     __snake_case     :					Union[str, Any]     ,     __snake_case     :					Optional[Any]     ,     __snake_case     :					Union[str, Any]     ,     __snake_case     :					Any     ,     __snake_case     :					str     ,     __snake_case     :					List[Any]    )							-> str:
            _lowerCAmelCase					=				self.num_choices
            _lowerCAmelCase					=				SqueezeBertForMultipleChoice(config=__snake_case    )
            model.to(__snake_case    )
            model.eval()
            _lowerCAmelCase					=				input_ids.unsqueeze(1    ).expand(-1     ,     self.num_choices     ,     -1    ).contiguous()
            _lowerCAmelCase					=				input_mask.unsqueeze(1    ).expand(-1     ,     self.num_choices     ,     -1    ).contiguous()
            _lowerCAmelCase					=				model(
                __snake_case     ,     attention_mask=__snake_case     ,     labels=__snake_case     ,     )
            self.parent.assertEqual(result.logits.shape     ,     (self.batch_size, self.num_choices)    )




      def 							lowercase__       (     self     :					Any    )							-> Tuple:
            _lowerCAmelCase					=				self.prepare_config_and_inputs()
            ((_lowerCAmelCase)	,				(_lowerCAmelCase)	,				(_lowerCAmelCase)	,				(_lowerCAmelCase)	,				(_lowerCAmelCase)	,				(_lowerCAmelCase))					=				config_and_inputs
            _lowerCAmelCase					=				{"""input_ids""": input_ids, """attention_mask""": input_mask}
            return config, inputs_dict




@require_torch
class 					UpperCAmelCase	(     snake_case_				,						snake_case_				,						unittest.TestCase  ):
      _lowercase:     Optional[int]			=  (
          (
              SqueezeBertModel,
              SqueezeBertForMaskedLM,
              SqueezeBertForMultipleChoice,
              SqueezeBertForQuestionAnswering,
              SqueezeBertForSequenceClassification,
              SqueezeBertForTokenClassification,
          )
          if is_torch_available()
          else None
      )
      _lowercase:     Optional[Any]			=  (
          {
              '''feature-extraction''': SqueezeBertModel,
              '''fill-mask''': SqueezeBertForMaskedLM,
              '''question-answering''': SqueezeBertForQuestionAnswering,
              '''text-classification''': SqueezeBertForSequenceClassification,
              '''token-classification''': SqueezeBertForTokenClassification,
              '''zero-shot''': SqueezeBertForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      _lowercase:     List[Any]			=  False
      _lowercase:     List[Any]			=  True
      _lowercase:     List[Any]			=  False




      def 							lowercase__       (     self     :					int    )							-> Dict:
            _lowerCAmelCase					=				SqueezeBertModelTester(self    )
            _lowerCAmelCase					=				ConfigTester(self     ,     config_class=__snake_case     ,     dim=37    )




      def 							lowercase__       (     self     :					List[Any]    )							-> List[Any]:
            self.config_tester.run_common_tests()




      def 							lowercase__       (     self     :					List[str]    )							-> Dict:
            _lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_squeezebert_model(*__snake_case    )




      def 							lowercase__       (     self     :					Dict    )							-> Optional[Any]:
            _lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_squeezebert_for_masked_lm(*__snake_case    )




      def 							lowercase__       (     self     :					Optional[Any]    )							-> Union[str, Any]:
            _lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_squeezebert_for_question_answering(*__snake_case    )




      def 							lowercase__       (     self     :					List[Any]    )							-> List[str]:
            _lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_squeezebert_for_sequence_classification(*__snake_case    )




      def 							lowercase__       (     self     :					Optional[int]    )							-> Any:
            _lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_squeezebert_for_token_classification(*__snake_case    )




      def 							lowercase__       (     self     :					Union[str, Any]    )							-> Tuple:
            _lowerCAmelCase					=				self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_squeezebert_for_multiple_choice(*__snake_case    )




      @slow
      def 							lowercase__       (     self     :					List[str]    )							-> Tuple:
            for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  _lowerCAmelCase					=				SqueezeBertModel.from_pretrained(__snake_case    )
                  self.assertIsNotNone(__snake_case    )




@require_sentencepiece
@require_tokenizers
@require_torch
class 					UpperCAmelCase	(     unittest.TestCase  ):




      @slow
      def 							lowercase__       (     self     :					Tuple    )							-> Union[str, Any]:
            _lowerCAmelCase					=				SqueezeBertForSequenceClassification.from_pretrained("""squeezebert/squeezebert-mnli"""    )

            _lowerCAmelCase					=				torch.tensor([[1, 2_94_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 13, 15_88, 2]]    )
            _lowerCAmelCase					=				model(__snake_case    )[0]
            _lowerCAmelCase					=				torch.Size((1, 3)    )
            self.assertEqual(output.shape     ,     __snake_case    )
            _lowerCAmelCase					=				torch.tensor([[0.64_01, -0.03_49, -0.60_41]]    )
            self.assertTrue(torch.allclose(__snake_case     ,     __snake_case     ,     atol=1E-4    )    )













'''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
)
A__							:  List[str]    =logging.getLogger(__name__)
def   UpperCamelCase__ (						):





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

      _lowerCAmelCase					=				[]
      _lowerCAmelCase					=				0
      _lowerCAmelCase					=				1_00_00
      _lowerCAmelCase					=				time.time()
      for text in data:
            _lowerCAmelCase					=				f"{bos} {text.strip()} {sep}"
            _lowerCAmelCase					=				tokenizer.encode(lowerCAmelCase				,      add_special_tokens=lowerCAmelCase    )
            rslt.append(lowerCAmelCase    )

            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(lowerCAmelCase    )} examples processed."    )

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


if __name__ == "__main__":
      main()






