




'''simple docstring'''






import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


lowercase	:      List[Any]				=       logging.getLogger(__name__)

require_version('''pytorch_lightning>=1.0.4''')

lowercase	:      Optional[int]				=       {
    '''base''': AutoModel,
    '''sequence-classification''': AutoModelForSequenceClassification,
    '''question-answering''': AutoModelForQuestionAnswering,
    '''pretraining''': AutoModelForPreTraining,
    '''token-classification''': AutoModelForTokenClassification,
    '''language-modeling''': AutoModelWithLMHead,
    '''summarization''': AutoModelForSeqaSeqLM,
    '''translation''': AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
lowercase	:      Any				=       {
    '''linear''': get_linear_schedule_with_warmup,
    '''cosine''': get_cosine_schedule_with_warmup,
    '''cosine_w_restarts''': get_cosine_with_hard_restarts_schedule_with_warmup,
    '''polynomial''': get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
lowercase	:      Any				=       sorted(arg_to_scheduler.keys())
lowercase	:      Tuple				=       '''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''

class   SCREAMING_SNAKE_CASE_		(       pl.LightningModule				):



 def __init__(  self   ,		A   ,		A=None   ,		A="base"   ,		A=None   ,		A=None   ,		A=None   ,		**A   ,		):






        '''simple docstring'''




        super().__init__()
        # TODO: move to self.save_hyperparameters()
        # self.save_hyperparameters()
        # can also expand arguments into trainer signature for easier reading

        self.save_hyperparameters(A )
        lowercase  :		Optional[Any]						=      0
        lowercase  :		Tuple						=      Path(self.hparams.output_dir )
        lowercase  :		Tuple						=      self.hparams.cache_dir if self.hparams.cache_dir else None
        if config is None:
               lowercase  :		List[str]						=      AutoConfig.from_pretrained(
                   self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path   ,		**({'''num_labels''': num_labels} if num_labels is not None else {})   ,		cache_dir=A   ,		**A   ,		)
        else:
               lowercase  :		PretrainedConfig						=      config

        lowercase  :		int						=      ('''encoder_layerdrop''', '''decoder_layerdrop''', '''dropout''', '''attention_dropout''')
        for p in extra_model_params:
               if getattr(self.hparams   ,		A   ,		A ):
                      assert hasattr(self.config   ,		A ), f"""model config doesn't have a `{p}` attribute"""
                      setattr(self.config   ,		A   ,		getattr(self.hparams   ,		A ) )

        if tokenizer is None:
               lowercase  :		List[str]						=      AutoTokenizer.from_pretrained(
                   self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path   ,		cache_dir=A   ,		)
        else:
               lowercase  :		PreTrainedTokenizer						=      tokenizer
        lowercase  :		Optional[Any]						=      MODEL_MODES[mode]
        if model is None:
               lowercase  :		List[Any]						=      self.model_type.from_pretrained(
                   self.hparams.model_name_or_path   ,		from_tf=bool('''.ckpt''' in self.hparams.model_name_or_path )   ,		config=self.config   ,		cache_dir=A   ,		)
        else:
               lowercase  :		List[str]						=      model



 def    SCREAMING_SNAKE_CASE(  self   ,		*A   ,		**A ):






        '''simple docstring'''




        lowercase  :		Optional[Any]						=      self.model_type.from_pretrained(*A   ,		**A )



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Dict						=      arg_to_scheduler[self.hparams.lr_scheduler]
        lowercase  :		List[Any]						=      get_schedule_func(
            self.opt   ,		num_warmup_steps=self.hparams.warmup_steps   ,		num_training_steps=self.total_steps() )
        lowercase  :		Tuple						=      {'''scheduler''': scheduler, '''interval''': '''step''', '''frequency''': 1}
        return scheduler



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Optional[Any]						=      self.model
        lowercase  :		Any						=      ['''bias''', '''LayerNorm.weight''']
        lowercase  :		Union[str, Any]						=      [
            {
                '''params''': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay )],
                '''weight_decay''': self.hparams.weight_decay,
            },
            {
                '''params''': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay )],
                '''weight_decay''': 0.0,
            },
        ]
        if self.hparams.adafactor:
               lowercase  :		str						=      Adafactor(
                   A   ,		lr=self.hparams.learning_rate   ,		scale_parameter=A   ,		relative_step=A )

        else:
               lowercase  :		str						=      AdamW(
                   A   ,		lr=self.hparams.learning_rate   ,		eps=self.hparams.adam_epsilon )
        lowercase  :		int						=      optimizer

        lowercase  :		Optional[Any]						=      self.get_lr_scheduler()

        return [optimizer], [scheduler]



 def    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A ):






        '''simple docstring'''




        return self.validation_step(A   ,		A )



 def    SCREAMING_SNAKE_CASE(  self   ,		A ):






        '''simple docstring'''




        return self.validation_end(A )



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Tuple						=      max(1   ,		self.hparams.gpus )  # TODO: consider num_tpu_cores
        lowercase  :		Optional[Any]						=      self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
        return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs



 def    SCREAMING_SNAKE_CASE(  self   ,		A ):






        '''simple docstring'''




        if stage == "test":
               lowercase  :		Tuple						=      len(self.test_dataloader().dataset )
        else:
               lowercase  :		Tuple						=      self.get_dataloader('''train'''   ,		self.hparams.train_batch_size   ,		shuffle=A )
               lowercase  :		List[str]						=      len(self.train_dataloader().dataset )



 def    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A   ,		A = False ):






        '''simple docstring'''




        raise NotImplementedError('''You must implement this for your task''' )



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        return self.train_loader



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        return self.get_dataloader('''dev'''   ,		self.hparams.eval_batch_size   ,		shuffle=A )



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        return self.get_dataloader('''test'''   ,		self.hparams.eval_batch_size   ,		shuffle=A )



 def    SCREAMING_SNAKE_CASE(  self   ,		A ):






        '''simple docstring'''




        return os.path.join(
            self.hparams.data_dir   ,		'''cached_{}_{}_{}'''.format(
                A   ,		list(filter(A   ,		self.hparams.model_name_or_path.split('''/''' ) ) ).pop()   ,		str(self.hparams.max_seq_length )   ,		)   ,		)



 @pl.utilities.rank_zero_only
 def    SCREAMING_SNAKE_CASE(  self   ,		A ):






        '''simple docstring'''




        lowercase  :		Optional[int]						=      self.output_dir.joinpath('''best_tfmr''' )
        lowercase  :		Optional[int]						=      self.step_count
        self.model.save_pretrained(A )
        self.tokenizer.save_pretrained(A )





 @staticmethod
 def    SCREAMING_SNAKE_CASE(  A   ,		A ):






        '''simple docstring'''




        parser.add_argument(
            '''--model_name_or_path'''   ,		default=A   ,		type=A   ,		required=A   ,		help='''Path to pretrained model or model identifier from huggingface.co/models'''   ,		)
        parser.add_argument(
            '''--config_name'''   ,		default=''''''   ,		type=A   ,		help='''Pretrained config name or path if not the same as model_name''' )
        parser.add_argument(
            '''--tokenizer_name'''   ,		default=A   ,		type=A   ,		help='''Pretrained tokenizer name or path if not the same as model_name'''   ,		)
        parser.add_argument(
            '''--cache_dir'''   ,		default=''''''   ,		type=A   ,		help='''Where do you want to store the pre-trained models downloaded from huggingface.co'''   ,		)
        parser.add_argument(
            '''--encoder_layerdrop'''   ,		type=A   ,		help='''Encoder layer dropout probability (Optional). Goes into model.config'''   ,		)
        parser.add_argument(
            '''--decoder_layerdrop'''   ,		type=A   ,		help='''Decoder layer dropout probability (Optional). Goes into model.config'''   ,		)
        parser.add_argument(
            '''--dropout'''   ,		type=A   ,		help='''Dropout probability (Optional). Goes into model.config'''   ,		)
        parser.add_argument(
            '''--attention_dropout'''   ,		type=A   ,		help='''Attention dropout probability (Optional). Goes into model.config'''   ,		)
        parser.add_argument('''--learning_rate'''   ,		default=5E-5   ,		type=A   ,		help='''The initial learning rate for Adam.''' )
        parser.add_argument(
            '''--lr_scheduler'''   ,		default='''linear'''   ,		choices=A   ,		metavar=A   ,		type=A   ,		help='''Learning rate scheduler'''   ,		)
        parser.add_argument('''--weight_decay'''   ,		default=0.0   ,		type=A   ,		help='''Weight decay if we apply some.''' )
        parser.add_argument('''--adam_epsilon'''   ,		default=1E-8   ,		type=A   ,		help='''Epsilon for Adam optimizer.''' )
        parser.add_argument('''--warmup_steps'''   ,		default=0   ,		type=A   ,		help='''Linear warmup over warmup_steps.''' )
        parser.add_argument('''--num_workers'''   ,		default=4   ,		type=A   ,		help='''kwarg passed to DataLoader''' )
        parser.add_argument('''--num_train_epochs'''   ,		dest='''max_epochs'''   ,		default=3   ,		type=A )
        parser.add_argument('''--train_batch_size'''   ,		default=32   ,		type=A )
        parser.add_argument('''--eval_batch_size'''   ,		default=32   ,		type=A )
        parser.add_argument('''--adafactor'''   ,		action='''store_true''' )






class   SCREAMING_SNAKE_CASE_		(       pl.Callback				):



 def    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A ):






        '''simple docstring'''




        if (
            trainer.is_global_zero and trainer.global_rank == 0
        ):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
               pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.






class   SCREAMING_SNAKE_CASE_		(       pl.Callback				):



 def    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A ):






        '''simple docstring'''




        lowercase  :		Any						=      trainer.lr_schedulers[0]['''scheduler''']
        lowercase  :		Tuple						=      {f"""lr_group_{i}""": lr for i, lr in enumerate(lr_scheduler.get_lr() )}
        pl_module.logger.log_metrics(A )



 def    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A ):






        '''simple docstring'''




        rank_zero_info('''***** Validation results *****''' )
        lowercase  :		List[Any]						=      trainer.callback_metrics
        # Log results
        for key in sorted(A ):
               if key not in ["log", "progress_bar"]:
                      rank_zero_info('''{} = {}\n'''.format(A   ,		str(metrics[key] ) ) )





 def    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A ):






        '''simple docstring'''




        rank_zero_info('''***** Test results *****''' )
        lowercase  :		Union[str, Any]						=      trainer.callback_metrics
        # Log and save results to file
        lowercase  :		Union[str, Any]						=      os.path.join(pl_module.hparams.output_dir   ,		'''test_results.txt''' )
        with open(A   ,		'''w''' ) as writer:
               for key in sorted(A ):
                      if key not in ["log", "progress_bar"]:
                             rank_zero_info('''{} = {}\n'''.format(A   ,		str(metrics[key] ) ) )
                             writer.write('''{} = {}\n'''.format(A   ,		str(metrics[key] ) ) )






def 		a_( lowerCamelCase   ,				lowerCamelCase      )				->None:






       '''simple docstring'''




       parser.add_argument(
           '''--output_dir'''   ,				default=lowerCamelCase   ,				type=lowerCamelCase   ,				required=lowerCamelCase   ,				help='''The output directory where the model predictions and checkpoints will be written.'''   ,				)
       parser.add_argument(
           '''--fp16'''   ,				action='''store_true'''   ,				help='''Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit'''   ,				)

       parser.add_argument(
           '''--fp16_opt_level'''   ,				type=lowerCamelCase   ,				default='''O2'''   ,				help=(
               '''For fp16: Apex AMP optimization level selected in [\'O0\', \'O1\', \'O2\', and \'O3\'].'''
               '''See details at https://nvidia.github.io/apex/amp.html'''
           )   ,				)
       parser.add_argument('''--n_tpu_cores'''   ,				dest='''tpu_cores'''   ,				type=lowerCamelCase      )
       parser.add_argument('''--max_grad_norm'''   ,				dest='''gradient_clip_val'''   ,				default=1.0   ,				type=lowerCamelCase   ,				help='''Max gradient norm'''      )
       parser.add_argument('''--do_train'''   ,				action='''store_true'''   ,				help='''Whether to run training.'''      )
       parser.add_argument('''--do_predict'''   ,				action='''store_true'''   ,				help='''Whether to run predictions on the test set.'''      )
       parser.add_argument(
           '''--gradient_accumulation_steps'''   ,				dest='''accumulate_grad_batches'''   ,				type=lowerCamelCase   ,				default=1   ,				help='''Number of updates steps to accumulate before performing a backward/update pass.'''   ,				)
       parser.add_argument('''--seed'''   ,				type=lowerCamelCase   ,				default=42   ,				help='''random seed for initialization'''      )
       parser.add_argument(
           '''--data_dir'''   ,				default=lowerCamelCase   ,				type=lowerCamelCase   ,				required=lowerCamelCase   ,				help='''The input data dir. Should contain the training files for the CoNLL-2003 NER task.'''   ,				)





def 		a_( lowerCamelCase   ,				lowerCamelCase   ,				lowerCamelCase=None   ,				lowerCamelCase=True   ,				lowerCamelCase=None   ,				lowerCamelCase=[]   ,				lowerCamelCase=None   ,				lowerCamelCase=None   ,				**lowerCamelCase   ,				)				->Dict:






       '''simple docstring'''




       pl.seed_everything(args.seed      )

       # init model
       lowercase  :		List[Any]						=      Path(model.hparams.output_dir      )
       odir.mkdir(exist_ok=lowerCamelCase      )

       # add custom checkpoints
       if checkpoint_callback is None:
              lowercase  :		Any						=      pl.callbacks.ModelCheckpoint(
                  filepath=args.output_dir   ,				prefix='''checkpoint'''   ,				monitor='''val_loss'''   ,				mode='''min'''   ,				save_top_k=1      )
       if early_stopping_callback:
              extra_callbacks.append(lowerCamelCase      )
       if logging_callback is None:
              lowercase  :		List[str]						=      LoggingCallback()

       lowercase  :		int						=      {}

       # TODO: remove with PyTorch 1.6 since pl uses native amp
       if args.fpaa:
              lowercase  :		Optional[Any]						=      16
              # train_params["amp_level"] = args.fp16_opt_level

       if args.gpus > 1:
              lowercase  :		Tuple						=      '''auto'''  # "ddp"
              lowercase  :		List[Any]						=      '''ddp'''

       lowercase  :		str						=      args.accumulate_grad_batches
       lowercase  :		Optional[Any]						=      None  # extra_train_kwargs.get("profiler", None) #get unwanted logs
       lowercase  :		Tuple						=      '''auto'''

       lowercase  :		Optional[Any]						=      pl.Trainer.from_argparse_args(
           lowerCamelCase   ,				weights_summary=lowerCamelCase   ,				callbacks=[logging_callback] + extra_callbacks + [checkpoint_callback] + [InitCallback()]   ,				logger=lowerCamelCase   ,				**lowerCamelCase   ,				)

       if args.do_train:
              trainer.fit(lowerCamelCase      )

       return trainer










'''simple docstring'''






import os
import pickle
import tempfile
import unittest
from typing import Callable, Optional

import numpy as np

from transformers import (
    BatchEncoding,
    BertTokenizer,
    BertTokenizerFast,
    PreTrainedTokenizer,
    PreTrainedTokenizerFast,
    TensorType,
    TokenSpan,
    is_tokenizers_available,
)
from transformers.models.gpta.tokenization_gpta import GPTaTokenizer
from transformers.testing_utils import CaptureStderr, require_flax, require_tf, require_tokenizers, require_torch, slow


if is_tokenizers_available():
      from tokenizers import Tokenizer
      from tokenizers.models import WordPiece

class   SCREAMING_SNAKE_CASE_		(       unittest.TestCase				):



 def    SCREAMING_SNAKE_CASE(  self   ,		A ):






        '''simple docstring'''




        lowercase  :		List[Any]						=      list(tokenizer_class.max_model_input_sizes.keys() )
        for model_name in sa_models[:1]:
               lowercase  :		Dict						=      tokenizer_class.from_pretrained(A )
               self.assertIsNotNone(A )
               self.assertIsInstance(A   ,		A )
               self.assertIsInstance(A   ,		A )

               for special_tok in tokenizer.all_special_tokens:
                      self.assertIsInstance(A   ,		A )
                      lowercase  :		Tuple						=      tokenizer.convert_tokens_to_ids(A )
                      self.assertIsInstance(A   ,		A )



 def    SCREAMING_SNAKE_CASE(  self   ,		A   ,		A = None ):






        '''simple docstring'''




        lowercase  :		List[Any]						=      pickle.dumps(A )
        self.assertIsNotNone(A )

        lowercase  :		str						=      pickle.loads(A )

        # Ensure is_fast is correctly restored
        self.assertEqual(be_restored.is_fast   ,		be_original.is_fast )

        # Ensure encodings are potentially correctly restored
        if be_original.is_fast:
               self.assertIsNotNone(be_restored.encodings )
        else:
               self.assertIsNone(be_restored.encodings )

        # Ensure the keys are the same
        for original_v, restored_v in zip(be_original.values()   ,		be_restored.values() ):
               if equal_op:
                      self.assertTrue(equal_op(A   ,		A ) )
               else:
                      self.assertEqual(A   ,		A )



 @slow
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        self.check_tokenizer_from_pretrained(A )



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        self.assertEqual(TensorType('''tf''' )   ,		TensorType.TENSORFLOW )
        self.assertEqual(TensorType('''pt''' )   ,		TensorType.PYTORCH )
        self.assertEqual(TensorType('''np''' )   ,		TensorType.NUMPY )



 @require_tokenizers
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        import numpy as np

        lowercase  :		Union[str, Any]						=      BertTokenizer.from_pretrained('''bert-base-cased''' )
        lowercase  :		str						=      BertTokenizerFast.from_pretrained('''bert-base-cased''' )

        # Python no tensor
        with self.subTest('''BatchEncoding (Python, return_tensors=None)''' ):
               self.assert_dump_and_restore(tokenizer_p('''Small example to encode''' ) )

        with self.subTest('''BatchEncoding (Python, return_tensors=NUMPY)''' ):
               self.assert_dump_and_restore(
                   tokenizer_p('''Small example to encode'''   ,		return_tensors=TensorType.NUMPY )   ,		np.array_equal )

        with self.subTest('''BatchEncoding (Rust, return_tensors=None)''' ):
               self.assert_dump_and_restore(tokenizer_r('''Small example to encode''' ) )

        with self.subTest('''BatchEncoding (Rust, return_tensors=NUMPY)''' ):
               self.assert_dump_and_restore(
                   tokenizer_r('''Small example to encode'''   ,		return_tensors=TensorType.NUMPY )   ,		np.array_equal )



 @require_tf
 @require_tokenizers
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        import tensorflow as tf

        def tf_array_equals(A   ,		A ):
               return tf.reduce_all(tf.equal(A   ,		A ) )

        lowercase  :		Optional[Any]						=      BertTokenizer.from_pretrained('''bert-base-cased''' )
        lowercase  :		Optional[Any]						=      BertTokenizerFast.from_pretrained('''bert-base-cased''' )

        with self.subTest('''BatchEncoding (Python, return_tensors=TENSORFLOW)''' ):
               self.assert_dump_and_restore(
                   tokenizer_p('''Small example to encode'''   ,		return_tensors=TensorType.TENSORFLOW )   ,		A )

        with self.subTest('''BatchEncoding (Rust, return_tensors=TENSORFLOW)''' ):
               self.assert_dump_and_restore(
                   tokenizer_r('''Small example to encode'''   ,		return_tensors=TensorType.TENSORFLOW )   ,		A )



 @require_torch
 @require_tokenizers
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        import torch

        lowercase  :		Tuple						=      BertTokenizer.from_pretrained('''bert-base-cased''' )
        lowercase  :		Any						=      BertTokenizerFast.from_pretrained('''bert-base-cased''' )

        with self.subTest('''BatchEncoding (Python, return_tensors=PYTORCH)''' ):
               self.assert_dump_and_restore(
                   tokenizer_p('''Small example to encode'''   ,		return_tensors=TensorType.PYTORCH )   ,		torch.equal )

        with self.subTest('''BatchEncoding (Rust, return_tensors=PYTORCH)''' ):
               self.assert_dump_and_restore(
                   tokenizer_r('''Small example to encode'''   ,		return_tensors=TensorType.PYTORCH )   ,		torch.equal )



 @require_tokenizers
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		str						=      BertTokenizer.from_pretrained('''bert-base-cased''' )
        lowercase  :		List[Any]						=      BertTokenizerFast.from_pretrained('''bert-base-cased''' )

        with self.subTest('''Python Tokenizer''' ):
               self.assertFalse(tokenizer_p('''Small example to_encode''' ).is_fast )

        with self.subTest('''Rust Tokenizer''' ):
               self.assertTrue(tokenizer_r('''Small example to_encode''' ).is_fast )



 @require_tokenizers
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Optional[int]						=      BertTokenizerFast.from_pretrained('''bert-base-cased''' )
        lowercase  :		str						=      tokenizer_r(['''Test''', '''\xad''', '''test''']   ,		is_split_into_words=A )

        self.assertEqual(encoded.word_to_tokens(0 )   ,		TokenSpan(start=1   ,		end=2 ) )
        self.assertEqual(encoded.word_to_tokens(1 )   ,		A )
        self.assertEqual(encoded.word_to_tokens(2 )   ,		TokenSpan(start=2   ,		end=3 ) )



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Optional[int]						=      BatchEncoding({'''inputs''': [[1, 2, 3], [4, 5, 6]], '''labels''': [0, 1]} )
        lowercase  :		List[str]						=      batch.convert_to_tensors(tensor_type='''np''' )
        self.assertEqual(tensor_batch['''inputs'''].shape   ,		(2, 3) )
        self.assertEqual(tensor_batch['''labels'''].shape   ,		(2,) )
        # test converting the converted
        with CaptureStderr() as cs:
               lowercase  :		Any						=      batch.convert_to_tensors(tensor_type='''np''' )
        self.assertFalse(len(cs.err )   ,		msg=f"""should have no warning, but got {cs.err}""" )

        lowercase  :		Dict						=      BatchEncoding({'''inputs''': [1, 2, 3], '''labels''': 0} )
        lowercase  :		List[str]						=      batch.convert_to_tensors(tensor_type='''np'''   ,		prepend_batch_axis=A )
        self.assertEqual(tensor_batch['''inputs'''].shape   ,		(1, 3) )
        self.assertEqual(tensor_batch['''labels'''].shape   ,		(1,) )



 @require_torch
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Dict						=      BatchEncoding({'''inputs''': [[1, 2, 3], [4, 5, 6]], '''labels''': [0, 1]} )
        lowercase  :		Union[str, Any]						=      batch.convert_to_tensors(tensor_type='''pt''' )
        self.assertEqual(tensor_batch['''inputs'''].shape   ,		(2, 3) )
        self.assertEqual(tensor_batch['''labels'''].shape   ,		(2,) )
        # test converting the converted
        with CaptureStderr() as cs:
               lowercase  :		Tuple						=      batch.convert_to_tensors(tensor_type='''pt''' )
        self.assertFalse(len(cs.err )   ,		msg=f"""should have no warning, but got {cs.err}""" )

        lowercase  :		Optional[Any]						=      BatchEncoding({'''inputs''': [1, 2, 3], '''labels''': 0} )
        lowercase  :		List[str]						=      batch.convert_to_tensors(tensor_type='''pt'''   ,		prepend_batch_axis=A )
        self.assertEqual(tensor_batch['''inputs'''].shape   ,		(1, 3) )
        self.assertEqual(tensor_batch['''labels'''].shape   ,		(1,) )



 @require_tf
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Dict						=      BatchEncoding({'''inputs''': [[1, 2, 3], [4, 5, 6]], '''labels''': [0, 1]} )
        lowercase  :		Dict						=      batch.convert_to_tensors(tensor_type='''tf''' )
        self.assertEqual(tensor_batch['''inputs'''].shape   ,		(2, 3) )
        self.assertEqual(tensor_batch['''labels'''].shape   ,		(2,) )
        # test converting the converted
        with CaptureStderr() as cs:
               lowercase  :		Optional[int]						=      batch.convert_to_tensors(tensor_type='''tf''' )
        self.assertFalse(len(cs.err )   ,		msg=f"""should have no warning, but got {cs.err}""" )

        lowercase  :		str						=      BatchEncoding({'''inputs''': [1, 2, 3], '''labels''': 0} )
        lowercase  :		List[str]						=      batch.convert_to_tensors(tensor_type='''tf'''   ,		prepend_batch_axis=A )
        self.assertEqual(tensor_batch['''inputs'''].shape   ,		(1, 3) )
        self.assertEqual(tensor_batch['''labels'''].shape   ,		(1,) )



 @require_flax
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Tuple						=      BatchEncoding({'''inputs''': [[1, 2, 3], [4, 5, 6]], '''labels''': [0, 1]} )
        lowercase  :		List[Any]						=      batch.convert_to_tensors(tensor_type='''jax''' )
        self.assertEqual(tensor_batch['''inputs'''].shape   ,		(2, 3) )
        self.assertEqual(tensor_batch['''labels'''].shape   ,		(2,) )
        # test converting the converted
        with CaptureStderr() as cs:
               lowercase  :		Dict						=      batch.convert_to_tensors(tensor_type='''jax''' )
        self.assertFalse(len(cs.err )   ,		msg=f"""should have no warning, but got {cs.err}""" )

        lowercase  :		str						=      BatchEncoding({'''inputs''': [1, 2, 3], '''labels''': 0} )
        lowercase  :		str						=      batch.convert_to_tensors(tensor_type='''jax'''   ,		prepend_batch_axis=A )
        self.assertEqual(tensor_batch['''inputs'''].shape   ,		(1, 3) )
        self.assertEqual(tensor_batch['''labels'''].shape   ,		(1,) )



 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		List[str]						=      [{'''input_ids''': np.array([0, 1, 2] )}, {'''input_ids''': np.array([0, 1, 2, 3] )}]
        lowercase  :		Tuple						=      BertTokenizer.from_pretrained('''bert-base-cased''' )

        lowercase  :		Dict						=      tokenizer.pad(A   ,		padding=A )
        self.assertTrue(isinstance(batch['''input_ids''']   ,		np.ndarray ) )
        self.assertEqual(batch['''input_ids'''].tolist()   ,		[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]] )
        lowercase  :		Dict						=      tokenizer.pad(A   ,		padding=A   ,		return_tensors='''np''' )
        self.assertTrue(isinstance(batch['''input_ids''']   ,		np.ndarray ) )
        self.assertEqual(batch['''input_ids'''].tolist()   ,		[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]] )



 @require_torch
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        import torch

        lowercase  :		Optional[Any]						=      [{'''input_ids''': torch.tensor([0, 1, 2] )}, {'''input_ids''': torch.tensor([0, 1, 2, 3] )}]
        lowercase  :		Tuple						=      BertTokenizer.from_pretrained('''bert-base-cased''' )

        lowercase  :		Any						=      tokenizer.pad(A   ,		padding=A )
        self.assertTrue(isinstance(batch['''input_ids''']   ,		torch.Tensor ) )
        self.assertEqual(batch['''input_ids'''].tolist()   ,		[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]] )
        lowercase  :		Union[str, Any]						=      tokenizer.pad(A   ,		padding=A   ,		return_tensors='''pt''' )
        self.assertTrue(isinstance(batch['''input_ids''']   ,		torch.Tensor ) )
        self.assertEqual(batch['''input_ids'''].tolist()   ,		[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]] )



 @require_tf
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        import tensorflow as tf

        lowercase  :		Optional[int]						=      [{'''input_ids''': tf.constant([0, 1, 2] )}, {'''input_ids''': tf.constant([0, 1, 2, 3] )}]
        lowercase  :		Dict						=      BertTokenizer.from_pretrained('''bert-base-cased''' )

        lowercase  :		List[Any]						=      tokenizer.pad(A   ,		padding=A )
        self.assertTrue(isinstance(batch['''input_ids''']   ,		tf.Tensor ) )
        self.assertEqual(batch['''input_ids'''].numpy().tolist()   ,		[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]] )
        lowercase  :		Any						=      tokenizer.pad(A   ,		padding=A   ,		return_tensors='''tf''' )
        self.assertTrue(isinstance(batch['''input_ids''']   ,		tf.Tensor ) )
        self.assertEqual(batch['''input_ids'''].numpy().tolist()   ,		[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]] )



 @require_tokenizers
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		Optional[Any]						=      Tokenizer(WordPiece(unk_token='''[UNK]''' ) )
        PreTrainedTokenizerFast(tokenizer_object=A )





 @require_tokenizers
 def    SCREAMING_SNAKE_CASE(  self ):






        '''simple docstring'''




        lowercase  :		List[str]						=      Tokenizer(WordPiece(unk_token='''[UNK]''' ) )
        with tempfile.TemporaryDirectory() as tmpdirname:
               bert_tokenizer.save(os.path.join(A   ,		'''tokenizer.json''' ) )
               PreTrainedTokenizerFast(tokenizer_file=os.path.join(A   ,		'''tokenizer.json''' ) )




