
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


_UpperCAmelCase      :				Tuple          = logging.getLogger(__name__)

require_version("pytorch_lightning>=1.0.4")

_UpperCAmelCase      :				Dict          = {
    "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
_UpperCAmelCase      :				Optional[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
}
_UpperCAmelCase      :				List[Any]          = sorted(arg_to_scheduler.keys())
_UpperCAmelCase      :				List[Any]          = "{" + ", ".join(arg_to_scheduler_choices) + "}"


class        lowercase      (	pl.LightningModule     ):



		def __init__(				self       ,   A_       ,   A_=None       ,   A_="base"       ,   A_=None       ,   A_=None       ,   A_=None       ,   **A_       ,   )						->      str:
				"""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_ )
				UpperCamelCase						=  0
				UpperCamelCase						=  Path(self.hparams.output_dir )
				UpperCamelCase						=  self.hparams.cache_dir if self.hparams.cache_dir else None
				if config is None:
						UpperCamelCase						=  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:
						UpperCamelCase						=  config

				UpperCamelCase						=  ('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:
						UpperCamelCase						=  AutoTokenizer.from_pretrained(
						    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path       ,   cache_dir=A_       ,   )
				else:
						UpperCamelCase						=  tokenizer
				UpperCamelCase						=  MODEL_MODES[mode]
				if model is None:
						UpperCamelCase						=  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:
						UpperCamelCase						=  model



		def __UpperCamelCase				(				self       ,   *A_       ,   **A_ )						->      List[str]:
				"""simple docstring"""
				UpperCamelCase						=  self.model_type.from_pretrained(*A_       ,   **A_ )



		def __UpperCamelCase				(				self )						->      Optional[int]:
				"""simple docstring"""
				UpperCamelCase						=  arg_to_scheduler[self.hparams.lr_scheduler]
				UpperCamelCase						=  get_schedule_func(
				    self.opt       ,   num_warmup_steps=self.hparams.warmup_steps       ,   num_training_steps=self.total_steps() )
				UpperCamelCase						=  {'scheduler': scheduler, 'interval': 'step', 'frequency': 1}
				return scheduler



		def __UpperCamelCase				(				self )						->      Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase						=  self.model
				UpperCamelCase						=  ['bias', 'LayerNorm.weight']
				UpperCamelCase						=  [
				    {
				        'params': [
				            p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay )
				        ],  # check this named paramters
				        '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:
						UpperCamelCase						=  Adafactor(
						    A_       ,   lr=self.hparams.learning_rate       ,   scale_parameter=A_       ,   relative_step=A_ )

				else:
						UpperCamelCase						=  AdamW(
						    A_       ,   lr=self.hparams.learning_rate       ,   eps=self.hparams.adam_epsilon )
				UpperCamelCase						=  optimizer

				UpperCamelCase						=  self.get_lr_scheduler()

				return [optimizer], [scheduler]



		def __UpperCamelCase				(				self       ,   A_       ,   A_ )						->      str:
				"""simple docstring"""
				return self.validation_step(A_       ,   A_ )



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[int]:
				"""simple docstring"""
				return self.validation_end(A_ )



		def __UpperCamelCase				(				self )						->      int:
				"""simple docstring"""
				UpperCamelCase						=  max(1       ,   self.hparams.gpus )  # TODO: consider num_tpu_cores
				UpperCamelCase						=  self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
				return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs



		def __UpperCamelCase				(				self       ,   A_ )						->      Union[str, Any]:
				"""simple docstring"""
				if stage == "test":
						UpperCamelCase						=  len(self.test_dataloader().dataset )
				else:
						UpperCamelCase						=  self.get_dataloader('train'       ,   self.hparams.train_batch_size       ,   shuffle=A_ )
						UpperCamelCase						=  len(self.train_dataloader().dataset )



		def __UpperCamelCase				(				self       ,   A_       ,   A_       ,   A_ = False )						->      List[str]:
				"""simple docstring"""
				raise NotImplementedError('You must implement this for your task' )



		def __UpperCamelCase				(				self )						->      Optional[int]:
				"""simple docstring"""
				return self.train_loader



		def __UpperCamelCase				(				self )						->      str:
				"""simple docstring"""
				return self.get_dataloader('dev'       ,   self.hparams.eval_batch_size       ,   shuffle=A_ )



		def __UpperCamelCase				(				self )						->      str:
				"""simple docstring"""
				return self.get_dataloader('test'       ,   self.hparams.eval_batch_size       ,   shuffle=A_ )



		def __UpperCamelCase				(				self       ,   A_ )						->      str:
				"""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 __UpperCamelCase				(				self       ,   A_ )						->      None:
				"""simple docstring"""
				UpperCamelCase						=  self.output_dir.joinpath('best_tfmr' )
				UpperCamelCase						=  self.step_count
				self.model.save_pretrained(A_ )
				self.tokenizer.save_pretrained(A_ )



		@staticmethod
		def __UpperCamelCase				(				A_       ,   A_ )						->      Any:
				"""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=str(Path(A_ ).parent / 'test_run' / 'cache' )       ,   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        lowercase      (	pl.Callback     ):



		def __UpperCamelCase				(				self       ,   A_       ,   A_ )						->      Optional[Any]:
				"""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        lowercase      (	pl.Callback     ):



		def __UpperCamelCase				(				self       ,   A_       ,   A_ )						->      Union[str, Any]:
				"""simple docstring"""
				# print(pl_module.model.rag)
				for name, param in pl_module.model.rag.named_parameters():
						if param.grad is None:
								print(A_ )






class        lowercase      (	pl.Callback     ):



		def __UpperCamelCase				(				self       ,   A_       ,   A_ )						->      Optional[int]:
				"""simple docstring"""
				UpperCamelCase						=  trainer.lr_schedulers[0]['scheduler']
				UpperCamelCase						=  {F'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr() )}
				pl_module.logger.log_metrics(A_ )



		def __UpperCamelCase				(				self       ,   A_       ,   A_ )						->      Dict:
				"""simple docstring"""
				rank_zero_info('***** Validation results *****' )
				UpperCamelCase						=  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 __UpperCamelCase				(				self       ,   A_       ,   A_ )						->      str:
				"""simple docstring"""
				rank_zero_info('***** Test results *****' )
				UpperCamelCase						=  trainer.callback_metrics
				# Log and save results to file
				UpperCamelCase						=  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 (      lowercase					, lowercase				)     ->						None:
		'''simple docstring'''

		parser.add_argument(
		    '--output_dir'					, default=str(Path(lowercase				).parent / 'test_run' / 'model_checkpoints'				)					, type=lowercase					, 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=lowercase					, 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=lowercase				)
		parser.add_argument('--max_grad_norm'					, dest='gradient_clip_val'					, default=1.0					, type=lowercase					, 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=lowercase					, default=1					, help='Number of updates steps to accumulate before performing a backward/update pass.'					, )
		parser.add_argument('--seed'					, type=lowercase					, default=42					, help='random seed for initialization'				)
		parser.add_argument(
		    '--data_dir'					, default=str(Path(lowercase				).parent / 'test_run' / 'dummy-train-data'				)					, type=lowercase					, help='The input data dir. Should contain the training files for the CoNLL-2003 NER task.'					, )


def    A (      lowercase					, lowercase					, lowercase=None					, lowercase=True					, lowercase=[]					, lowercase=None					, lowercase=None					, **lowercase					, )     ->						Dict:
		'''simple docstring'''

		pl.seed_everything(args.seed				)

		# init model
		UpperCamelCase						=  Path(model.hparams.output_dir				)
		odir.mkdir(exist_ok=lowercase				)

		# add custom checkpoints
		if checkpoint_callback is None:
				UpperCamelCase						=  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(lowercase				)
		if logging_callback is None:
				UpperCamelCase						=  LoggingCallback()

		UpperCamelCase						=  {}

		if args.fpaa:
				UpperCamelCase						=  16

		if args.gpus > 1:
				UpperCamelCase						=  'auto'
				UpperCamelCase						=  'ddp'

		UpperCamelCase						=  args.accumulate_grad_batches
		UpperCamelCase						=  None
		UpperCamelCase						=  'auto'

		UpperCamelCase						=  pl.Trainer.from_argparse_args(
		    lowercase					, weights_summary=lowercase					, callbacks=[logging_callback] + extra_callbacks + [InitCallback()] + [checkpoint_callback]					, logger=lowercase					, val_check_interval=1					, num_sanity_val_steps=2					, **lowercase					, )

		if args.do_train:
				trainer.fit(lowercase				)

		else:
				print('RAG modeling tests with new set functions successfuly executed!'				)
		return trainer


import os
import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from huggingface_hub.file_download import http_get
from requests.exceptions import HTTPError

from transformers import (
    AlbertTokenizer,
    AutoTokenizer,
    BertTokenizer,
    BertTokenizerFast,
    GPTaTokenizerFast,
    is_tokenizers_available,
)
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_tokenizers
from transformers.tokenization_utils import Trie


sys.path.append(str(Path(__file__).parent.parent / "utils"))

from test_module.custom_tokenization import CustomTokenizer  # noqa E402


if is_tokenizers_available():
					from test_module.custom_tokenization_fast import CustomTokenizerFast


class        lowercase      (	unittest.TestCase     ):



		def __UpperCamelCase				(				self )						->      Tuple:
				"""simple docstring"""
				# A mock response for an HTTP head request to emulate server down
				UpperCamelCase						=  mock.Mock()
				UpperCamelCase						=  500
				UpperCamelCase						=  {}
				UpperCamelCase						=  HTTPError
				UpperCamelCase						=  {}

				# Download this model to make sure it's in the cache.
				UpperCamelCase						=  BertTokenizer.from_pretrained('hf-internal-testing/tiny-random-bert' )

				# Under the mock environment we get a 500 error when trying to reach the tokenizer.
				with mock.patch('requests.Session.request'       ,   return_value=A_ ) as mock_head:
						UpperCamelCase						=  BertTokenizer.from_pretrained('hf-internal-testing/tiny-random-bert' )
						# This check we did call the fake head request
						mock_head.assert_called()



		@require_tokenizers
		def __UpperCamelCase				(				self )						->      Dict:
				"""simple docstring"""
				# A mock response for an HTTP head request to emulate server down
				UpperCamelCase						=  mock.Mock()
				UpperCamelCase						=  500
				UpperCamelCase						=  {}
				UpperCamelCase						=  HTTPError
				UpperCamelCase						=  {}

				# Download this model to make sure it's in the cache.
				UpperCamelCase						=  GPTaTokenizerFast.from_pretrained('gpt2' )

				# Under the mock environment we get a 500 error when trying to reach the tokenizer.
				with mock.patch('requests.Session.request'       ,   return_value=A_ ) as mock_head:
						UpperCamelCase						=  GPTaTokenizerFast.from_pretrained('gpt2' )
						# This check we did call the fake head request
						mock_head.assert_called()



		def __UpperCamelCase				(				self )						->      Optional[Any]:
				"""simple docstring"""
				# This test is for deprecated behavior and can be removed in v5
				try:
						UpperCamelCase						=  tempfile.mktemp()
						with open(A_       ,   'wb' ) as f:
								http_get('https://huggingface.co/albert-base-v1/resolve/main/spiece.model'       ,   A_ )

						UpperCamelCase						=  AlbertTokenizer.from_pretrained(A_ )
				finally:
						os.remove(A_ )

				# Supporting this legacy load introduced a weird bug where the tokenizer would load local files if they are in
				# the current folder and have the right name.
				if os.path.isfile('tokenizer.json' ):
						# We skip the test if the user has a `tokenizer.json` in this folder to avoid deleting it.
						return
				try:
						with open('tokenizer.json'       ,   'wb' ) as f:
								http_get('https://huggingface.co/hf-internal-testing/tiny-random-bert/blob/main/tokenizer.json'       ,   A_ )
						UpperCamelCase						=  AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2' )
						# The tiny random BERT has a vocab size of 1024, tiny gpt2 as a vocab size of 1000
						self.assertEqual(tokenizer.vocab_size       ,   1_000 )
						# Tokenizer should depend on the remote checkpoint, not the local tokenizer.json file.

				finally:
						os.remove('tokenizer.json' )



		def __UpperCamelCase				(				self )						->      Dict:
				"""simple docstring"""
				# This test is for deprecated behavior and can be removed in v5
				UpperCamelCase						=  AlbertTokenizer.from_pretrained('https://huggingface.co/albert-base-v1/resolve/main/spiece.model' )






@is_staging_test
class        lowercase      (	unittest.TestCase     ):
		__lowercase : int			  =       ["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]", "bla", "blou"]



		@classmethod
		def __UpperCamelCase				(				cls )						->      Tuple:
				"""simple docstring"""
				UpperCamelCase						=  TOKEN
				HfFolder.save_token(A_ )



		@classmethod
		def __UpperCamelCase				(				cls )						->      Optional[int]:
				"""simple docstring"""
				try:
						delete_repo(token=cls._token       ,   repo_id='test-tokenizer' )
				except HTTPError:
						pass

				try:
						delete_repo(token=cls._token       ,   repo_id='valid_org/test-tokenizer-org' )
				except HTTPError:
						pass

				try:
						delete_repo(token=cls._token       ,   repo_id='test-dynamic-tokenizer' )
				except HTTPError:
						pass



		def __UpperCamelCase				(				self )						->      Any:
				"""simple docstring"""
				with tempfile.TemporaryDirectory() as tmp_dir:
						UpperCamelCase						=  os.path.join(A_       ,   'vocab.txt' )
						with open(A_       ,   'w'       ,   encoding='utf-8' ) as vocab_writer:
								vocab_writer.write(''.join([x + '\n' for x in self.vocab_tokens] ) )
						UpperCamelCase						=  BertTokenizer(A_ )

				tokenizer.push_to_hub('test-tokenizer'       ,   use_auth_token=self._token )
				UpperCamelCase						=  BertTokenizer.from_pretrained(F'''{USER}/test-tokenizer''' )
				self.assertDictEqual(new_tokenizer.vocab       ,   tokenizer.vocab )

				# Reset repo
				delete_repo(token=self._token       ,   repo_id='test-tokenizer' )

				# Push to hub via save_pretrained
				with tempfile.TemporaryDirectory() as tmp_dir:
						tokenizer.save_pretrained(A_       ,   repo_id='test-tokenizer'       ,   push_to_hub=A_       ,   use_auth_token=self._token )

				UpperCamelCase						=  BertTokenizer.from_pretrained(F'''{USER}/test-tokenizer''' )
				self.assertDictEqual(new_tokenizer.vocab       ,   tokenizer.vocab )



		def __UpperCamelCase				(				self )						->      int:
				"""simple docstring"""
				with tempfile.TemporaryDirectory() as tmp_dir:
						UpperCamelCase						=  os.path.join(A_       ,   'vocab.txt' )
						with open(A_       ,   'w'       ,   encoding='utf-8' ) as vocab_writer:
								vocab_writer.write(''.join([x + '\n' for x in self.vocab_tokens] ) )
						UpperCamelCase						=  BertTokenizer(A_ )

				tokenizer.push_to_hub('valid_org/test-tokenizer-org'       ,   use_auth_token=self._token )
				UpperCamelCase						=  BertTokenizer.from_pretrained('valid_org/test-tokenizer-org' )
				self.assertDictEqual(new_tokenizer.vocab       ,   tokenizer.vocab )

				# Reset repo
				delete_repo(token=self._token       ,   repo_id='valid_org/test-tokenizer-org' )

				# Push to hub via save_pretrained
				with tempfile.TemporaryDirectory() as tmp_dir:
						tokenizer.save_pretrained(
						    A_       ,   repo_id='valid_org/test-tokenizer-org'       ,   push_to_hub=A_       ,   use_auth_token=self._token )

				UpperCamelCase						=  BertTokenizer.from_pretrained('valid_org/test-tokenizer-org' )
				self.assertDictEqual(new_tokenizer.vocab       ,   tokenizer.vocab )



		@require_tokenizers
		def __UpperCamelCase				(				self )						->      Dict:
				"""simple docstring"""
				CustomTokenizer.register_for_auto_class()
				with tempfile.TemporaryDirectory() as tmp_dir:
						UpperCamelCase						=  os.path.join(A_       ,   'vocab.txt' )
						with open(A_       ,   'w'       ,   encoding='utf-8' ) as vocab_writer:
								vocab_writer.write(''.join([x + '\n' for x in self.vocab_tokens] ) )
						UpperCamelCase						=  CustomTokenizer(A_ )

				# No fast custom tokenizer
				tokenizer.push_to_hub('test-dynamic-tokenizer'       ,   use_auth_token=self._token )

				UpperCamelCase						=  AutoTokenizer.from_pretrained(F'''{USER}/test-dynamic-tokenizer'''       ,   trust_remote_code=A_ )
				# Can't make an isinstance check because the new_model.config is from the CustomTokenizer class of a dynamic module
				self.assertEqual(tokenizer.__class__.__name__       ,   'CustomTokenizer' )

				# Fast and slow custom tokenizer
				CustomTokenizerFast.register_for_auto_class()
				with tempfile.TemporaryDirectory() as tmp_dir:
						UpperCamelCase						=  os.path.join(A_       ,   'vocab.txt' )
						with open(A_       ,   'w'       ,   encoding='utf-8' ) as vocab_writer:
								vocab_writer.write(''.join([x + '\n' for x in self.vocab_tokens] ) )

						UpperCamelCase						=  BertTokenizerFast.from_pretrained(A_ )
						bert_tokenizer.save_pretrained(A_ )
						UpperCamelCase						=  CustomTokenizerFast.from_pretrained(A_ )

				tokenizer.push_to_hub('test-dynamic-tokenizer'       ,   use_auth_token=self._token )

				UpperCamelCase						=  AutoTokenizer.from_pretrained(F'''{USER}/test-dynamic-tokenizer'''       ,   trust_remote_code=A_ )
				# Can't make an isinstance check because the new_model.config is from the FakeConfig class of a dynamic module
				self.assertEqual(tokenizer.__class__.__name__       ,   'CustomTokenizerFast' )
				UpperCamelCase						=  AutoTokenizer.from_pretrained(
				    F'''{USER}/test-dynamic-tokenizer'''       ,   use_fast=A_       ,   trust_remote_code=A_ )
				# Can't make an isinstance check because the new_model.config is from the FakeConfig class of a dynamic module
				self.assertEqual(tokenizer.__class__.__name__       ,   'CustomTokenizer' )






class        lowercase      (	unittest.TestCase     ):



		def __UpperCamelCase				(				self )						->      Optional[int]:
				"""simple docstring"""
				UpperCamelCase						=  Trie()
				trie.add('Hello 友達' )
				self.assertEqual(trie.data       ,   {'H': {'e': {'l': {'l': {'o': {' ': {'友': {'達': {'': 1}}}}}}}}} )
				trie.add('Hello' )
				trie.data
				self.assertEqual(trie.data       ,   {'H': {'e': {'l': {'l': {'o': {'': 1, ' ': {'友': {'達': {'': 1}}}}}}}}} )



		def __UpperCamelCase				(				self )						->      str:
				"""simple docstring"""
				UpperCamelCase						=  Trie()
				self.assertEqual(trie.split('[CLS] This is a extra_id_100' )       ,   ['[CLS] This is a extra_id_100'] )
				trie.add('[CLS]' )
				trie.add('extra_id_1' )
				trie.add('extra_id_100' )
				self.assertEqual(trie.split('[CLS] This is a extra_id_100' )       ,   ['[CLS]', ' This is a ', 'extra_id_100'] )



		def __UpperCamelCase				(				self )						->      List[Any]:
				"""simple docstring"""
				UpperCamelCase						=  Trie()
				trie.add('A' )
				self.assertEqual(trie.split('ABC' )       ,   ['A', 'BC'] )
				self.assertEqual(trie.split('BCA' )       ,   ['BC', 'A'] )



		def __UpperCamelCase				(				self )						->      Tuple:
				"""simple docstring"""
				UpperCamelCase						=  Trie()
				trie.add('TOKEN]' )
				trie.add('[SPECIAL_TOKEN]' )
				self.assertEqual(trie.split('This is something [SPECIAL_TOKEN]' )       ,   ['This is something ', '[SPECIAL_TOKEN]'] )



		def __UpperCamelCase				(				self )						->      Dict:
				"""simple docstring"""
				UpperCamelCase						=  Trie()
				trie.add('A' )
				trie.add('P' )
				trie.add('[SPECIAL_TOKEN]' )
				self.assertEqual(trie.split('This is something [SPECIAL_TOKEN]' )       ,   ['This is something ', '[SPECIAL_TOKEN]'] )



		def __UpperCamelCase				(				self )						->      Optional[Any]:
				"""simple docstring"""
				UpperCamelCase						=  Trie()
				trie.add('AB' )
				trie.add('B' )
				trie.add('C' )
				self.assertEqual(trie.split('ABC' )       ,   ['AB', 'C'] )



		def __UpperCamelCase				(				self )						->      Tuple:
				"""simple docstring"""
				UpperCamelCase						=  Trie()
				trie.add('ABC' )
				trie.add('B' )
				trie.add('CD' )
				self.assertEqual(trie.split('ABCD' )       ,   ['ABC', 'D'] )



		def __UpperCamelCase				(				self )						->      int:
				"""simple docstring"""
				# Even if the offsets are wrong, we necessarily output correct string
				# parts.
				UpperCamelCase						=  Trie()
				UpperCamelCase						=  trie.cut_text('ABC'       ,   [0, 0, 2, 1, 2, 3] )
				self.assertEqual(A_       ,   ['AB', 'C'] )
