






import os
import shutil
import sys
import tempfile
import unittest
from pathlib import Path

import pytest

import transformers
from transformers import (
    BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP,
    AutoTokenizer,
    BertConfig,
    BertTokenizer,
    BertTokenizerFast,
    CTRLTokenizer,
    GPTaTokenizer,
    GPTaTokenizerFast,
    PreTrainedTokenizerFast,
    RobertaTokenizer,
    RobertaTokenizerFast,
    is_tokenizers_available,
)
from transformers.models.auto.configuration_auto import CONFIG_MAPPING, AutoConfig
from transformers.models.auto.tokenization_auto import (
    TOKENIZER_MAPPING,
    get_tokenizer_config,
    tokenizer_class_from_name,
)
from transformers.models.roberta.configuration_roberta import RobertaConfig
from transformers.testing_utils import (
    DUMMY_DIFF_TOKENIZER_IDENTIFIER,
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tokenizers,
    slow,
)


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

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_tokenization import CustomTokenizer  # noqa E402


if is_tokenizers_available():
      from test_module.custom_tokenization_fast import CustomTokenizerFast







class 				UpperCAmelCase_  (		unittest.TestCase		):




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  Union[str, Any]				      =			0




     @slow
     def   __UpperCAmelCase  (				self		):
          for model_name in (x for x in BERT_PRETRAINED_CONFIG_ARCHIVE_MAP.keys() if "japanese" not in x):
               UpperCAmelCase__      :  Union[str, Any]				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase		)
               self.assertIsNotNone(_lowerCAmelCase		)
               self.assertIsInstance(_lowerCAmelCase	,						(BertTokenizer, BertTokenizerFast)		)
               self.assertGreater(len(_lowerCAmelCase		)	,						0		)

          for model_name in GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP.keys():
               UpperCAmelCase__      :  Any				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase		)
               self.assertIsNotNone(_lowerCAmelCase		)
               self.assertIsInstance(_lowerCAmelCase	,						(GPTaTokenizer, GPTaTokenizerFast)		)
               self.assertGreater(len(_lowerCAmelCase		)	,						0		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  Optional[Any]				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase		)
          self.assertIsInstance(_lowerCAmelCase	,						(BertTokenizer, BertTokenizerFast)		)
          self.assertEqual(tokenizer.vocab_size	,						12		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  str				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase		)
          self.assertIsInstance(_lowerCAmelCase	,						(RobertaTokenizer, RobertaTokenizerFast)		)
          self.assertEqual(tokenizer.vocab_size	,						20		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  int				      =			AutoConfig.from_pretrained(_lowerCAmelCase		)
          self.assertIsInstance(_lowerCAmelCase	,						_lowerCAmelCase		)
          # Check that tokenizer_type ≠ model_type
          UpperCAmelCase__      :  Dict				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase	,						config=_lowerCAmelCase		)
          self.assertIsInstance(_lowerCAmelCase	,						(BertTokenizer, BertTokenizerFast)		)
          self.assertEqual(tokenizer.vocab_size	,						12		)




     def   __UpperCAmelCase  (				self		):
          with tempfile.TemporaryDirectory() as tmp_dir:
               shutil.copy("""./tests/fixtures/vocab.txt"""	,						os.path.join(_lowerCAmelCase	,						"""vocab.txt"""		)		)

               UpperCAmelCase__      :  Any				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase	,						tokenizer_type="""bert"""	,						use_fast=_lowerCAmelCase		)
               self.assertIsInstance(_lowerCAmelCase	,						_lowerCAmelCase		)

          with tempfile.TemporaryDirectory() as tmp_dir:
               shutil.copy("""./tests/fixtures/vocab.json"""	,						os.path.join(_lowerCAmelCase	,						"""vocab.json"""		)		)
               shutil.copy("""./tests/fixtures/merges.txt"""	,						os.path.join(_lowerCAmelCase	,						"""merges.txt"""		)		)

               UpperCAmelCase__      :  Union[str, Any]				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase	,						tokenizer_type="""gpt2"""	,						use_fast=_lowerCAmelCase		)
               self.assertIsInstance(_lowerCAmelCase	,						_lowerCAmelCase		)




     @require_tokenizers
     def   __UpperCAmelCase  (				self		):
          with tempfile.TemporaryDirectory() as tmp_dir:
               shutil.copy("""./tests/fixtures/vocab.txt"""	,						os.path.join(_lowerCAmelCase	,						"""vocab.txt"""		)		)

               UpperCAmelCase__      :  int				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase	,						tokenizer_type="""bert"""		)
               self.assertIsInstance(_lowerCAmelCase	,						_lowerCAmelCase		)

          with tempfile.TemporaryDirectory() as tmp_dir:
               shutil.copy("""./tests/fixtures/vocab.json"""	,						os.path.join(_lowerCAmelCase	,						"""vocab.json"""		)		)
               shutil.copy("""./tests/fixtures/merges.txt"""	,						os.path.join(_lowerCAmelCase	,						"""merges.txt"""		)		)

               UpperCAmelCase__      :  Any				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase	,						tokenizer_type="""gpt2"""		)
               self.assertIsInstance(_lowerCAmelCase	,						_lowerCAmelCase		)




     def   __UpperCAmelCase  (				self		):
          with pytest.raises(_lowerCAmelCase		):
               AutoTokenizer.from_pretrained("""./"""	,						tokenizer_type="""xxx"""		)




     @require_tokenizers
     def   __UpperCAmelCase  (				self		):
          for tokenizer_class in [BertTokenizer, BertTokenizerFast, AutoTokenizer]:
               UpperCAmelCase__      :  Tuple				      =			tokenizer_class.from_pretrained("""wietsedv/bert-base-dutch-cased"""		)
               self.assertIsInstance(_lowerCAmelCase	,						(BertTokenizer, BertTokenizerFast)		)

               if isinstance(_lowerCAmelCase	,						_lowerCAmelCase		):
                    self.assertEqual(tokenizer.basic_tokenizer.do_lower_case	,						_lowerCAmelCase		)
               else:
                    self.assertEqual(tokenizer.do_lower_case	,						_lowerCAmelCase		)

               self.assertEqual(tokenizer.model_max_length	,						512		)




     @require_tokenizers
     def   __UpperCAmelCase  (				self		):
          for tokenizer_class in [BertTokenizer, BertTokenizerFast, AutoTokenizer]:
               with self.assertRaisesRegex(
                   _lowerCAmelCase	,						"""julien-c/herlolip-not-exists is not a local folder and is not a valid model identifier"""	,						):
                    UpperCAmelCase__      :  Optional[int]				      =			tokenizer_class.from_pretrained("""julien-c/herlolip-not-exists"""		)




     def   __UpperCAmelCase  (				self		):
          # tests: https://github.com/huggingface/transformers/pull/13251
          # 1. models with `-`, e.g. xlm-roberta -> xlm_roberta
          # 2. models that don't remap 1-1 from model-name to model file, e.g., openai-gpt -> openai
          UpperCAmelCase__      :  List[str]				      =			TOKENIZER_MAPPING.values()
          UpperCAmelCase__      :  Dict				      =			[]

          for slow_tok, fast_tok in tokenizers:
               if slow_tok is not None:
                    tokenizer_names.append(slow_tok.__name__		)

               if fast_tok is not None:
                    tokenizer_names.append(fast_tok.__name__		)

          for tokenizer_name in tokenizer_names:
               # must find the right class
               tokenizer_class_from_name(_lowerCAmelCase		)




     @require_tokenizers
     def   __UpperCAmelCase  (				self		):
          self.assertIsInstance(AutoTokenizer.from_pretrained("""bert-base-cased"""	,						use_fast=_lowerCAmelCase		)	,						_lowerCAmelCase		)
          self.assertIsInstance(AutoTokenizer.from_pretrained("""bert-base-cased"""		)	,						_lowerCAmelCase		)




     @require_tokenizers
     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  str				      =			AutoTokenizer.from_pretrained("""distilbert-base-uncased"""	,						do_lower_case=_lowerCAmelCase		)
          UpperCAmelCase__      :  Any				      =			"""Hello, world. How are you?"""
          UpperCAmelCase__      :  Any				      =			tokenizer.tokenize(_lowerCAmelCase		)
          self.assertEqual("""[UNK]"""	,						tokens[0]		)

          UpperCAmelCase__      :  Dict				      =			AutoTokenizer.from_pretrained("""microsoft/mpnet-base"""	,						do_lower_case=_lowerCAmelCase		)
          UpperCAmelCase__      :  Optional[int]				      =			tokenizer.tokenize(_lowerCAmelCase		)
          self.assertEqual("""[UNK]"""	,						tokens[0]		)




     @require_tokenizers
     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  int				      =			AutoTokenizer.from_pretrained("""robot-test/dummy-tokenizer-fast-with-model-config"""		)
          self.assertEqual(type(_lowerCAmelCase		)	,						_lowerCAmelCase		)
          self.assertEqual(tokenizer.model_max_length	,						512		)
          self.assertEqual(tokenizer.vocab_size	,						30000		)
          self.assertEqual(tokenizer.unk_token	,						"""[UNK]"""		)
          self.assertEqual(tokenizer.padding_side	,						"""right"""		)
          self.assertEqual(tokenizer.truncation_side	,						"""right"""		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  Union[str, Any]				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase		)
          self.assertIsInstance(_lowerCAmelCase	,						(BertTokenizer, BertTokenizerFast)		)
          with tempfile.TemporaryDirectory() as tmp_dir:
               tokenizer.save_pretrained(_lowerCAmelCase		)
               UpperCAmelCase__      :  str				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase		)

          self.assertIsInstance(_lowerCAmelCase	,						tokenizer.__class__		)
          self.assertEqual(tokenizera.vocab_size	,						12		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  int				      =			AutoTokenizer.from_pretrained("""ctrl"""		)
          # There is no fast CTRL so this always gives us a slow tokenizer.
          self.assertIsInstance(_lowerCAmelCase	,						_lowerCAmelCase		)




     def   __UpperCAmelCase  (				self		):
          # Check we can load the tokenizer config of an online model.
          UpperCAmelCase__      :  Union[str, Any]				      =			get_tokenizer_config("""bert-base-cased"""		)
          UpperCAmelCase__      :  List[str]				      =			config.pop("""_commit_hash"""	,						_lowerCAmelCase		)
          # If we ever update bert-base-cased tokenizer config, this dict here will need to be updated.
          self.assertEqual(_lowerCAmelCase	,						{"""do_lower_case""": False}		)

          # This model does not have a tokenizer_config so we get back an empty dict.
          UpperCAmelCase__      :  Union[str, Any]				      =			get_tokenizer_config(_lowerCAmelCase		)
          self.assertDictEqual(_lowerCAmelCase	,						{}		)

          # A tokenizer saved with `save_pretrained` always creates a tokenizer config.
          UpperCAmelCase__      :  Union[str, Any]				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase		)
          with tempfile.TemporaryDirectory() as tmp_dir:
               tokenizer.save_pretrained(_lowerCAmelCase		)
               UpperCAmelCase__      :  Any				      =			get_tokenizer_config(_lowerCAmelCase		)

          # Check the class of the tokenizer was properly saved (note that it always saves the slow class).
          self.assertEqual(config["""tokenizer_class"""]	,						"""BertTokenizer"""		)




     def   __UpperCAmelCase  (				self		):
          try:
               AutoConfig.register("""custom"""	,						_lowerCAmelCase		)

               AutoTokenizer.register(_lowerCAmelCase	,						slow_tokenizer_class=_lowerCAmelCase		)
               # Trying to register something existing in the Transformers library will raise an error
               with self.assertRaises(_lowerCAmelCase		):
                    AutoTokenizer.register(_lowerCAmelCase	,						slow_tokenizer_class=_lowerCAmelCase		)

               UpperCAmelCase__      :  Optional[Any]				      =			CustomTokenizer.from_pretrained(_lowerCAmelCase		)
               with tempfile.TemporaryDirectory() as tmp_dir:
                    tokenizer.save_pretrained(_lowerCAmelCase		)

                    UpperCAmelCase__      :  Any				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase		)
                    self.assertIsInstance(_lowerCAmelCase	,						_lowerCAmelCase		)

          finally:
               if "custom" in CONFIG_MAPPING._extra_content:
                    del CONFIG_MAPPING._extra_content["custom"]
               if CustomConfig in TOKENIZER_MAPPING._extra_content:
                    del TOKENIZER_MAPPING._extra_content[CustomConfig]




     @require_tokenizers
     def   __UpperCAmelCase  (				self		):
          try:
               AutoConfig.register("""custom"""	,						_lowerCAmelCase		)

               # Can register in two steps
               AutoTokenizer.register(_lowerCAmelCase	,						slow_tokenizer_class=_lowerCAmelCase		)
               self.assertEqual(TOKENIZER_MAPPING[CustomConfig]	,						(CustomTokenizer, None)		)
               AutoTokenizer.register(_lowerCAmelCase	,						fast_tokenizer_class=_lowerCAmelCase		)
               self.assertEqual(TOKENIZER_MAPPING[CustomConfig]	,						(CustomTokenizer, CustomTokenizerFast)		)

               del TOKENIZER_MAPPING._extra_content[CustomConfig]
               # Can register in one step
               AutoTokenizer.register(
                   _lowerCAmelCase	,						slow_tokenizer_class=_lowerCAmelCase	,						fast_tokenizer_class=_lowerCAmelCase		)
               self.assertEqual(TOKENIZER_MAPPING[CustomConfig]	,						(CustomTokenizer, CustomTokenizerFast)		)

               # Trying to register something existing in the Transformers library will raise an error
               with self.assertRaises(_lowerCAmelCase		):
                    AutoTokenizer.register(_lowerCAmelCase	,						fast_tokenizer_class=_lowerCAmelCase		)

               # We pass through a bert tokenizer fast cause there is no converter slow to fast for our new toknizer
               # and that model does not have a tokenizer.json
               with tempfile.TemporaryDirectory() as tmp_dir:
                    UpperCAmelCase__      :  Dict				      =			BertTokenizerFast.from_pretrained(_lowerCAmelCase		)
                    bert_tokenizer.save_pretrained(_lowerCAmelCase		)
                    UpperCAmelCase__      :  Union[str, Any]				      =			CustomTokenizerFast.from_pretrained(_lowerCAmelCase		)

               with tempfile.TemporaryDirectory() as tmp_dir:
                    tokenizer.save_pretrained(_lowerCAmelCase		)

                    UpperCAmelCase__      :  Union[str, Any]				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase		)
                    self.assertIsInstance(_lowerCAmelCase	,						_lowerCAmelCase		)

                    UpperCAmelCase__      :  List[Any]				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase	,						use_fast=_lowerCAmelCase		)
                    self.assertIsInstance(_lowerCAmelCase	,						_lowerCAmelCase		)

          finally:
               if "custom" in CONFIG_MAPPING._extra_content:
                    del CONFIG_MAPPING._extra_content["custom"]
               if CustomConfig in TOKENIZER_MAPPING._extra_content:
                    del TOKENIZER_MAPPING._extra_content[CustomConfig]




     def   __UpperCAmelCase  (				self		):
          # If remote code is not set, we will time out when asking whether to load the model.
          with self.assertRaises(_lowerCAmelCase		):
               UpperCAmelCase__      :  int				      =			AutoTokenizer.from_pretrained("""hf-internal-testing/test_dynamic_tokenizer"""		)
          # If remote code is disabled, we can't load this config.
          with self.assertRaises(_lowerCAmelCase		):
               UpperCAmelCase__      :  List[str]				      =			AutoTokenizer.from_pretrained(
                   """hf-internal-testing/test_dynamic_tokenizer"""	,						trust_remote_code=_lowerCAmelCase		)

          UpperCAmelCase__      :  List[str]				      =			AutoTokenizer.from_pretrained("""hf-internal-testing/test_dynamic_tokenizer"""	,						trust_remote_code=_lowerCAmelCase		)
          self.assertTrue(tokenizer.special_attribute_present		)
          # Test tokenizer can be reloaded.
          with tempfile.TemporaryDirectory() as tmp_dir:
               tokenizer.save_pretrained(_lowerCAmelCase		)
               UpperCAmelCase__      :  Any				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase	,						trust_remote_code=_lowerCAmelCase		)
          self.assertTrue(reloaded_tokenizer.special_attribute_present		)

          if is_tokenizers_available():
               self.assertEqual(tokenizer.__class__.__name__	,						"""NewTokenizerFast"""		)
               self.assertEqual(reloaded_tokenizer.__class__.__name__	,						"""NewTokenizerFast"""		)

               # Test we can also load the slow version
               UpperCAmelCase__      :  Optional[Any]				      =			AutoTokenizer.from_pretrained(
                   """hf-internal-testing/test_dynamic_tokenizer"""	,						trust_remote_code=_lowerCAmelCase	,						use_fast=_lowerCAmelCase		)
               self.assertTrue(tokenizer.special_attribute_present		)
               self.assertEqual(tokenizer.__class__.__name__	,						"""NewTokenizer"""		)
               # Test tokenizer can be reloaded.
               with tempfile.TemporaryDirectory() as tmp_dir:
                    tokenizer.save_pretrained(_lowerCAmelCase		)
                    UpperCAmelCase__      :  Union[str, Any]				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase	,						trust_remote_code=_lowerCAmelCase	,						use_fast=_lowerCAmelCase		)
               self.assertEqual(reloaded_tokenizer.__class__.__name__	,						"""NewTokenizer"""		)
               self.assertTrue(reloaded_tokenizer.special_attribute_present		)
          else:
               self.assertEqual(tokenizer.__class__.__name__	,						"""NewTokenizer"""		)
               self.assertEqual(reloaded_tokenizer.__class__.__name__	,						"""NewTokenizer"""		)




     @require_tokenizers
     def   __UpperCAmelCase  (				self		):







          class 				UpperCAmelCase_  (		__lowerCamelCase		):
               __lowerCamelCase      				=  False







          class 				UpperCAmelCase_  (		__lowerCamelCase		):
               __lowerCamelCase      				=  NewTokenizer
               __lowerCamelCase      				=  False

          try:
               AutoConfig.register("""custom"""	,						_lowerCAmelCase		)
               AutoTokenizer.register(_lowerCAmelCase	,						slow_tokenizer_class=_lowerCAmelCase		)
               AutoTokenizer.register(_lowerCAmelCase	,						fast_tokenizer_class=_lowerCAmelCase		)
               # If remote code is not set, the default is to use local
               UpperCAmelCase__      :  Any				      =			AutoTokenizer.from_pretrained("""hf-internal-testing/test_dynamic_tokenizer"""		)
               self.assertEqual(tokenizer.__class__.__name__	,						"""NewTokenizerFast"""		)
               self.assertFalse(tokenizer.special_attribute_present		)
               UpperCAmelCase__      :  int				      =			AutoTokenizer.from_pretrained("""hf-internal-testing/test_dynamic_tokenizer"""	,						use_fast=_lowerCAmelCase		)
               self.assertEqual(tokenizer.__class__.__name__	,						"""NewTokenizer"""		)
               self.assertFalse(tokenizer.special_attribute_present		)

               # If remote code is disabled, we load the local one.
               UpperCAmelCase__      :  Dict				      =			AutoTokenizer.from_pretrained(
                   """hf-internal-testing/test_dynamic_tokenizer"""	,						trust_remote_code=_lowerCAmelCase		)
               self.assertEqual(tokenizer.__class__.__name__	,						"""NewTokenizerFast"""		)
               self.assertFalse(tokenizer.special_attribute_present		)
               UpperCAmelCase__      :  Tuple				      =			AutoTokenizer.from_pretrained(
                   """hf-internal-testing/test_dynamic_tokenizer"""	,						trust_remote_code=_lowerCAmelCase	,						use_fast=_lowerCAmelCase		)
               self.assertEqual(tokenizer.__class__.__name__	,						"""NewTokenizer"""		)
               self.assertFalse(tokenizer.special_attribute_present		)

               # If remote is enabled, we load from the Hub
               UpperCAmelCase__      :  List[Any]				      =			AutoTokenizer.from_pretrained(
                   """hf-internal-testing/test_dynamic_tokenizer"""	,						trust_remote_code=_lowerCAmelCase		)
               self.assertEqual(tokenizer.__class__.__name__	,						"""NewTokenizerFast"""		)
               self.assertTrue(tokenizer.special_attribute_present		)
               UpperCAmelCase__      :  List[Any]				      =			AutoTokenizer.from_pretrained(
                   """hf-internal-testing/test_dynamic_tokenizer"""	,						trust_remote_code=_lowerCAmelCase	,						use_fast=_lowerCAmelCase		)
               self.assertEqual(tokenizer.__class__.__name__	,						"""NewTokenizer"""		)
               self.assertTrue(tokenizer.special_attribute_present		)

          finally:
               if "custom" in CONFIG_MAPPING._extra_content:
                    del CONFIG_MAPPING._extra_content["custom"]
               if CustomConfig in TOKENIZER_MAPPING._extra_content:
                    del TOKENIZER_MAPPING._extra_content[CustomConfig]




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  int				      =			AutoTokenizer.from_pretrained(
              """hf-internal-testing/test_dynamic_tokenizer_legacy"""	,						trust_remote_code=_lowerCAmelCase		)
          self.assertTrue(tokenizer.special_attribute_present		)
          if is_tokenizers_available():
               self.assertEqual(tokenizer.__class__.__name__	,						"""NewTokenizerFast"""		)

               # Test we can also load the slow version
               UpperCAmelCase__      :  List[Any]				      =			AutoTokenizer.from_pretrained(
                   """hf-internal-testing/test_dynamic_tokenizer_legacy"""	,						trust_remote_code=_lowerCAmelCase	,						use_fast=_lowerCAmelCase		)
               self.assertTrue(tokenizer.special_attribute_present		)
               self.assertEqual(tokenizer.__class__.__name__	,						"""NewTokenizer"""		)
          else:
               self.assertEqual(tokenizer.__class__.__name__	,						"""NewTokenizer"""		)




     def   __UpperCAmelCase  (				self		):
          with self.assertRaisesRegex(
              _lowerCAmelCase	,						"""bert-base is not a local folder and is not a valid model identifier"""		):
               UpperCAmelCase__      :  Any				      =			AutoTokenizer.from_pretrained("""bert-base"""		)




     def   __UpperCAmelCase  (				self		):
          with self.assertRaisesRegex(
              _lowerCAmelCase	,						R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""		):
               UpperCAmelCase__      :  Union[str, Any]				      =			AutoTokenizer.from_pretrained(_lowerCAmelCase	,						revision="""aaaaaa"""		)




     def   __UpperCAmelCase  (				self		):
          # Make sure we have cached the tokenizer.
          UpperCAmelCase__      :  Optional[Any]				      =			AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-bert"""		)
          with RequestCounter() as counter:
               UpperCAmelCase__      :  Any				      =			AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-bert"""		)
               self.assertEqual(counter.get_request_count	,						0		)
               self.assertEqual(counter.head_request_count	,						1		)
               self.assertEqual(counter.other_request_count	,						0		)














"""simple docstring"""
import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
			import jax
			import jax.numpy as jnp
			from flax.jax_utils import replicate
			from flax.training.common_utils import shard


@slow
@require_flax
class _lowerCAmelCase						(    unittest.TestCase       ):







						def 			A      (    self							)		->							Any:
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()







						def 			A      (    self							)		->							int:
										_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    FlaxStableDiffusionPipeline.from_pretrained(
										    'stabilityai/stable-diffusion-2'   ,				revision='bf16'   ,				dtype=jnp.bfloataa   ,				)

										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    'A painting of a squirrel eating a burger'
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    jax.device_count()
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    num_samples * [prompt]
										_SCREAMING_SNAKE_CASE   :		Any       		=    sd_pipe.prepare_inputs(lowerCAmelCase_							)

										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    replicate(lowerCAmelCase_							)
										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    shard(lowerCAmelCase_							)

										_SCREAMING_SNAKE_CASE   :		Any       		=    jax.random.PRNGKey(0							)
										_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    jax.random.split(lowerCAmelCase_   ,				jax.device_count()							)

										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    sd_pipe(lowerCAmelCase_   ,				lowerCAmelCase_   ,				lowerCAmelCase_   ,				num_inference_steps=2_5   ,				jit=lowerCAmelCase_							)[0]
										assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

										_SCREAMING_SNAKE_CASE   :		int       		=    images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]							)
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

										_SCREAMING_SNAKE_CASE   :		Tuple       		=    jnp.asarray(jax.device_get(image_slice.flatten()							)							)
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    jnp.array([0.4_238, 0.4_414, 0.4_395, 0.4_453, 0.4_629, 0.4_590, 0.4_531, 0.45_508, 0.4_512]							)
										print(F"""output_slice: {output_slice}"""							)
										assert jnp.abs(output_slice - expected_slice							).max() < 1e-2







						def 			A      (    self							)		->							Optional[int]:
										_SCREAMING_SNAKE_CASE   :		str       		=    'stabilityai/stable-diffusion-2'
										_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE   :		int       		=    FlaxDPMSolverMultistepScheduler.from_pretrained(lowerCAmelCase_   ,				subfolder='scheduler'							)
										_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE   :		List[str]       		=    FlaxStableDiffusionPipeline.from_pretrained(
										    lowerCAmelCase_   ,				scheduler=lowerCAmelCase_   ,				revision='bf16'   ,				dtype=jnp.bfloataa   ,				)
										_SCREAMING_SNAKE_CASE   :		Any       		=    scheduler_params

										_SCREAMING_SNAKE_CASE   :		int       		=    'A painting of a squirrel eating a burger'
										_SCREAMING_SNAKE_CASE   :		Tuple       		=    jax.device_count()
										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    num_samples * [prompt]
										_SCREAMING_SNAKE_CASE   :		Any       		=    sd_pipe.prepare_inputs(lowerCAmelCase_							)

										_SCREAMING_SNAKE_CASE   :		Dict       		=    replicate(lowerCAmelCase_							)
										_SCREAMING_SNAKE_CASE   :		Any       		=    shard(lowerCAmelCase_							)

										_SCREAMING_SNAKE_CASE   :		Any       		=    jax.random.PRNGKey(0							)
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    jax.random.split(lowerCAmelCase_   ,				jax.device_count()							)

										_SCREAMING_SNAKE_CASE   :		int       		=    sd_pipe(lowerCAmelCase_   ,				lowerCAmelCase_   ,				lowerCAmelCase_   ,				num_inference_steps=2_5   ,				jit=lowerCAmelCase_							)[0]
										assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

										_SCREAMING_SNAKE_CASE   :		str       		=    images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]							)
										_SCREAMING_SNAKE_CASE   :		str       		=    images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

										_SCREAMING_SNAKE_CASE   :		int       		=    jnp.asarray(jax.device_get(image_slice.flatten()							)							)
										_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    jnp.array([0.4_336, 0.42_969, 0.4_453, 0.4_199, 0.4_297, 0.4_531, 0.4_434, 0.4_434, 0.4_297]							)
										print(F"""output_slice: {output_slice}"""							)
										assert jnp.abs(output_slice - expected_slice							).max() < 1e-2

