






import json
import os
import shutil
import tempfile
import unittest

from transformers import BatchEncoding, CanineTokenizer
from transformers.testing_utils import require_tokenizers, require_torch
from transformers.tokenization_utils import AddedToken
from transformers.utils import cached_property

from ...test_tokenization_common import TokenizerTesterMixin





class 			__a  (    a_       ,   unittest.TestCase       ):
 _lowerCAmelCase       :       Optional[int]        = CanineTokenizer
 _lowerCAmelCase       :       Union[str, Any]        = False







 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






      super().setUp()
      UpperCamelCase__			:   Dict       				=  CanineTokenizer()
      tokenizer.save_pretrained(self.tmpdirname						)







 @cached_property
 def 					__lowercase					(		self					:   str						):


      '''simple docstring'''






      return CanineTokenizer.from_pretrained("google/canine-s"						)







 def 					__lowercase					(		self					:   Optional[Any]							,   **SCREAMING_SNAKE_CASE					:   List[Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  self.tokenizer_class.from_pretrained(self.tmpdirname							,   **_lowerCamelCase						)
      UpperCamelCase__			:   Dict       				=  10_24
      return tokenizer







 @require_torch
 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  self.canine_tokenizer
      UpperCamelCase__			:   str       				=  ["Life is like a box of chocolates.", "You never know what you\'re gonna get."]
      # fmt: off
      UpperCamelCase__			:   str       				=  [5_73_44, 76, 1_05, 1_02, 1_01, 32, 1_05, 1_15, 32, 1_08, 1_05, 1_07, 1_01, 32, 97, 32, 98, 1_11, 1_20, 32, 1_11, 1_02, 32, 99, 1_04, 1_11, 99, 1_11, 1_08, 97, 1_16, 1_01, 1_15, 46, 5_73_45, 0, 0, 0, 0]
      # fmt: on
      UpperCamelCase__			:   str       				=  tokenizer(_lowerCamelCase							,   padding=_lowerCamelCase							,   return_tensors="pt"						)
      self.assertIsInstance(_lowerCamelCase							,   _lowerCamelCase						)

      UpperCamelCase__			:   List[str]       				=  list(batch.input_ids.numpy()[0]						)

      self.assertListEqual(_lowerCamelCase							,   _lowerCamelCase						)

      self.assertEqual((2, 39)							,   batch.input_ids.shape						)
      self.assertEqual((2, 39)							,   batch.attention_mask.shape						)







 @require_torch
 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[Any]       				=  self.canine_tokenizer
      UpperCamelCase__			:   List[str]       				=  ["Once there was a man.", "He wrote a test in HuggingFace Tranformers."]
      UpperCamelCase__			:   Optional[Any]       				=  tokenizer(_lowerCamelCase							,   padding=_lowerCamelCase							,   return_tensors="pt"						)
      # check if input_ids, attention_mask and token_type_ids are returned
      self.assertIn("input_ids"							,   _lowerCamelCase						)
      self.assertIn("attention_mask"							,   _lowerCamelCase						)
      self.assertIn("token_type_ids"							,   _lowerCamelCase						)







 @require_torch
 def 					__lowercase					(		self					:   Union[str, Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Dict       				=  self.canine_tokenizer
      UpperCamelCase__			:   Dict       				=  [
          "What\'s the weater?",
          "It\'s about 25 degrees.",
      ]
      UpperCamelCase__			:   str       				=  tokenizer(
          text_target=_lowerCamelCase							,   max_length=32							,   padding="max_length"							,   truncation=_lowerCamelCase							,   return_tensors="pt"						)
      self.assertEqual(32							,   targets["input_ids"].shape[1]						)







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  self.get_tokenizers()
      for tokenizer in tokenizers:
           with self.subTest(F'{tokenizer.__class__.__name__}'						):
                self.assertNotEqual(tokenizer.model_max_length							,   42						)

           # Now let's start the test
      UpperCamelCase__			:   Optional[int]       				=  self.get_tokenizers()
      for tokenizer in tokenizers:
           with self.subTest(F'{tokenizer.__class__.__name__}'						):
                # Isolate this from the other tests because we save additional tokens/etc
                UpperCamelCase__			:   str       				=  tempfile.mkdtemp()

                UpperCamelCase__			:   int       				=  " He is very happy, UNwant\u00E9d,running"
                UpperCamelCase__			:   List[Any]       				=  tokenizer.encode(_lowerCamelCase							,   add_special_tokens=_lowerCamelCase						)
                tokenizer.save_pretrained(_lowerCamelCase						)

                UpperCamelCase__			:   List[Any]       				=  tokenizer.__class__.from_pretrained(_lowerCamelCase						)
                UpperCamelCase__			:   Optional[Any]       				=  after_tokenizer.encode(_lowerCamelCase							,   add_special_tokens=_lowerCamelCase						)
                self.assertListEqual(_lowerCamelCase							,   _lowerCamelCase						)

                shutil.rmtree(_lowerCamelCase						)

      UpperCamelCase__			:   Union[str, Any]       				=  self.get_tokenizers(model_max_length=42						)
      for tokenizer in tokenizers:
           with self.subTest(F'{tokenizer.__class__.__name__}'						):
                # Isolate this from the other tests because we save additional tokens/etc
                UpperCamelCase__			:   List[Any]       				=  tempfile.mkdtemp()

                UpperCamelCase__			:   Optional[int]       				=  " He is very happy, UNwant\u00E9d,running"

                UpperCamelCase__			:   Union[str, Any]       				=  tokenizer.additional_special_tokens

                # We can add a new special token for Canine as follows:
                UpperCamelCase__			:   Dict       				=  chr(0Xe0_07						)
                additional_special_tokens.append(_lowerCamelCase						)
                tokenizer.add_special_tokens({"additional_special_tokens": additional_special_tokens}						)
                UpperCamelCase__			:   Optional[Any]       				=  tokenizer.encode(_lowerCamelCase							,   add_special_tokens=_lowerCamelCase						)
                tokenizer.save_pretrained(_lowerCamelCase						)

                UpperCamelCase__			:   Dict       				=  tokenizer.__class__.from_pretrained(_lowerCamelCase						)
                UpperCamelCase__			:   List[Any]       				=  after_tokenizer.encode(_lowerCamelCase							,   add_special_tokens=_lowerCamelCase						)
                self.assertListEqual(_lowerCamelCase							,   _lowerCamelCase						)
                self.assertIn(_lowerCamelCase							,   after_tokenizer.additional_special_tokens						)
                self.assertEqual(after_tokenizer.model_max_length							,   42						)

                UpperCamelCase__			:   List[str]       				=  tokenizer.__class__.from_pretrained(_lowerCamelCase							,   model_max_length=43						)
                self.assertEqual(tokenizer.model_max_length							,   43						)

                shutil.rmtree(_lowerCamelCase						)







 def 					__lowercase					(		self					:   str						):


      '''simple docstring'''






      UpperCamelCase__			:   Union[str, Any]       				=  self.get_tokenizers(do_lower_case=_lowerCamelCase						)
      for tokenizer in tokenizers:
           with self.subTest(F'{tokenizer.__class__.__name__}'						):
                UpperCamelCase__				,     UpperCamelCase__			:   str       				=  self.get_clean_sequence(_lowerCamelCase						)

                # a special token for Canine can be defined as follows:
                UpperCamelCase__			:   Optional[int]       				=  0Xe0_05
                UpperCamelCase__			:   str       				=  chr(_lowerCamelCase						)

                tokenizer.add_special_tokens({"cls_token": special_token}						)
                UpperCamelCase__			:   Tuple       				=  tokenizer.encode(_lowerCamelCase							,   add_special_tokens=_lowerCamelCase						)
                self.assertEqual(len(_lowerCamelCase						)							,   1						)

                UpperCamelCase__			:   str       				=  tokenizer.decode(ids + encoded_special_token							,   clean_up_tokenization_spaces=_lowerCamelCase						)
                UpperCamelCase__			:   Optional[int]       				=  tokenizer.encode(_lowerCamelCase							,   add_special_tokens=_lowerCamelCase						)

                UpperCamelCase__			:   Optional[int]       				=  tokenizer.encode(_lowerCamelCase							,   add_special_tokens=_lowerCamelCase						)
                UpperCamelCase__			:   List[str]       				=  tokenizer.encode(_lowerCamelCase							,   add_special_tokens=_lowerCamelCase						)
                self.assertEqual(_lowerCamelCase							,   input_encoded + special_token_id						)

                UpperCamelCase__			:   str       				=  tokenizer.decode(_lowerCamelCase							,   skip_special_tokens=_lowerCamelCase						)
                self.assertTrue(special_token not in decoded						)







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   Union[str, Any]       				=  self.get_tokenizers(do_lower_case=_lowerCamelCase						)
      for tokenizer in tokenizers:
           with self.subTest(F'{tokenizer.__class__.__name__}'						):
                UpperCamelCase__			:   Tuple       				=  chr(0Xe0_05						)
                UpperCamelCase__			:   int       				=  chr(0Xe0_06						)

                # `add_tokens` method stores special tokens only in `tokenizer.unique_no_split_tokens`. (in tokenization_utils.py)
                tokenizer.add_tokens([SPECIAL_TOKEN_1]							,   special_tokens=_lowerCamelCase						)
                # `add_special_tokens` method stores special tokens in `tokenizer.additional_special_tokens`,
                # which also occur in `tokenizer.all_special_tokens`. (in tokenization_utils_base.py)
                tokenizer.add_special_tokens({"additional_special_tokens": [SPECIAL_TOKEN_2]}						)

                UpperCamelCase__			:   Optional[Any]       				=  tokenizer.tokenize(_lowerCamelCase						)
                UpperCamelCase__			:   List[str]       				=  tokenizer.tokenize(_lowerCamelCase						)

                self.assertEqual(len(_lowerCamelCase						)							,   1						)
                self.assertEqual(len(_lowerCamelCase						)							,   1						)
                self.assertEqual(token_a[0]							,   _lowerCamelCase						)
                self.assertEqual(token_a[0]							,   _lowerCamelCase						)







 @require_tokenizers
 def 					__lowercase					(		self					:   Optional[Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  self.get_tokenizers(do_lower_case=_lowerCamelCase						)
      for tokenizer in tokenizers:
           with self.subTest(F'{tokenizer.__class__.__name__}'						):
                # a special token for Canine can be defined as follows:
                UpperCamelCase__			:   int       				=  0Xe0_06
                UpperCamelCase__			:   List[str]       				=  chr(_lowerCamelCase						)

                UpperCamelCase__			:   Dict       				=  AddedToken(_lowerCamelCase							,   lstrip=_lowerCamelCase						)
                tokenizer.add_special_tokens({"additional_special_tokens": [new_token]}						)

                with tempfile.TemporaryDirectory() as tmp_dir_name:
                     tokenizer.save_pretrained(_lowerCamelCase						)
                     tokenizer.from_pretrained(_lowerCamelCase						)







 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  []
      if self.test_slow_tokenizer:
           tokenizer_list.append((self.tokenizer_class, self.get_tokenizer())						)

      if self.test_rust_tokenizer:
           tokenizer_list.append((self.rust_tokenizer_class, self.get_rust_tokenizer())						)

      for tokenizer_class, tokenizer_utils in tokenizer_list:
           with tempfile.TemporaryDirectory() as tmp_dir:
                tokenizer_utils.save_pretrained(_lowerCamelCase						)

                with open(os.path.join(_lowerCamelCase							,   "special_tokens_map.json"						)							,   encoding="utf-8"						) as json_file:
                     UpperCamelCase__			:   int       				=  json.load(_lowerCamelCase						)

                with open(os.path.join(_lowerCamelCase							,   "tokenizer_config.json"						)							,   encoding="utf-8"						) as json_file:
                     UpperCamelCase__			:   Optional[Any]       				=  json.load(_lowerCamelCase						)

                # a special token for Canine can be defined as follows:
                UpperCamelCase__			:   Dict       				=  0Xe0_06
                UpperCamelCase__			:   List[Any]       				=  chr(_lowerCamelCase						)

                UpperCamelCase__			:   Union[str, Any]       				=  [new_token_a]
                UpperCamelCase__			:   List[str]       				=  [new_token_a]

                with open(os.path.join(_lowerCamelCase							,   "special_tokens_map.json"						)							,   "w"							,   encoding="utf-8"						) as outfile:
                     json.dump(_lowerCamelCase							,   _lowerCamelCase						)
                with open(os.path.join(_lowerCamelCase							,   "tokenizer_config.json"						)							,   "w"							,   encoding="utf-8"						) as outfile:
                     json.dump(_lowerCamelCase							,   _lowerCamelCase						)

                # the following checks allow us to verify that our test works as expected, i.e. that the tokenizer takes
                # into account the new value of additional_special_tokens given in the "tokenizer_config.json" and
                # "special_tokens_map.json" files
                UpperCamelCase__			:   Tuple       				=  tokenizer_class.from_pretrained(_lowerCamelCase							,   extra_ids=0						)
                self.assertIn(_lowerCamelCase							,   tokenizer_without_change_in_init.additional_special_tokens						)
                # self.assertIn("an_additional_special_token",tokenizer_without_change_in_init.get_vocab()) # ByT5Tokenization no vocab
                self.assertEqual(
                    [new_token_a]							,   tokenizer_without_change_in_init.convert_ids_to_tokens(
                        tokenizer_without_change_in_init.convert_tokens_to_ids([new_token_a]						)						)							,   )

                UpperCamelCase__			:   Any       				=  0Xe0_07
                UpperCamelCase__			:   Optional[int]       				=  chr(_lowerCamelCase						)
                # Now we test that we can change the value of additional_special_tokens in the from_pretrained
                UpperCamelCase__			:   List[str]       				=  [AddedToken(_lowerCamelCase							,   lstrip=_lowerCamelCase						)]
                UpperCamelCase__			:   Optional[Any]       				=  tokenizer_class.from_pretrained(
                    _lowerCamelCase							,   additional_special_tokens=_lowerCamelCase							,   extra_ids=0						)

                self.assertIn(_lowerCamelCase							,   tokenizer.additional_special_tokens						)
                # self.assertIn(new_token_2,tokenizer.get_vocab()) # ByT5Tokenization no vocab
                self.assertEqual(
                    [new_token_a]							,   tokenizer.convert_ids_to_tokens(tokenizer.convert_tokens_to_ids([new_token_a]						)						)						)







 @require_tokenizers
 def 					__lowercase					(		self					:   Dict						):


      '''simple docstring'''






      UpperCamelCase__			:   Union[str, Any]       				=  self.get_tokenizers(do_lower_case=_lowerCamelCase						)
      for tokenizer in tokenizers:
           with self.subTest(F'{tokenizer.__class__.__name__}'						):
                UpperCamelCase__			:   List[str]       				=  "hello world"
                if self.space_between_special_tokens:
                     UpperCamelCase__			:   str       				=  "[CLS] hello world [SEP]"
                else:
                     UpperCamelCase__			:   List[str]       				=  input
                UpperCamelCase__			:   Union[str, Any]       				=  tokenizer.encode(_lowerCamelCase							,   add_special_tokens=_lowerCamelCase						)
                UpperCamelCase__			:   Dict       				=  tokenizer.decode(_lowerCamelCase							,   spaces_between_special_tokens=self.space_between_special_tokens						)
                self.assertIn(_lowerCamelCase							,   [output, output.lower()]						)







 def 					__lowercase					(		self					:   Union[str, Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Union[str, Any]       				=  self.get_tokenizers()
      for tokenizer in tokenizers:
           with self.subTest(F'{tokenizer.__class__.__name__}'						):
                UpperCamelCase__			:   Union[str, Any]       				=  [
                    "bos_token",
                    "eos_token",
                    "unk_token",
                    "sep_token",
                    "pad_token",
                    "cls_token",
                    "mask_token",
                ]

      UpperCamelCase__			:   str       				=  "a"
      UpperCamelCase__			:   Union[str, Any]       				=  ord(_lowerCamelCase						)

      for attr in attributes_list:
           setattr(_lowerCamelCase							,   attr + "_id"							,   _lowerCamelCase						)
           self.assertEqual(getattr(_lowerCamelCase							,   _lowerCamelCase						)							,   _lowerCamelCase						)
           self.assertEqual(getattr(_lowerCamelCase							,   attr + "_id"						)							,   _lowerCamelCase						)

           setattr(_lowerCamelCase							,   attr + "_id"							,   _lowerCamelCase						)
           self.assertEqual(getattr(_lowerCamelCase							,   _lowerCamelCase						)							,   _lowerCamelCase						)
           self.assertEqual(getattr(_lowerCamelCase							,   attr + "_id"						)							,   _lowerCamelCase						)

      setattr(_lowerCamelCase							,   "additional_special_tokens_ids"							,   []						)
      self.assertListEqual(getattr(_lowerCamelCase							,   "additional_special_tokens"						)							,   []						)
      self.assertListEqual(getattr(_lowerCamelCase							,   "additional_special_tokens_ids"						)							,   []						)

      UpperCamelCase__			:   List[str]       				=  0Xe0_06
      UpperCamelCase__			:   Dict       				=  chr(_lowerCamelCase						)
      setattr(_lowerCamelCase							,   "additional_special_tokens_ids"							,   [additional_special_token_id]						)
      self.assertListEqual(getattr(_lowerCamelCase							,   "additional_special_tokens"						)							,   [additional_special_token]						)
      self.assertListEqual(getattr(_lowerCamelCase							,   "additional_special_tokens_ids"						)							,   [additional_special_token_id]						)







 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      pass







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      pass







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      pass







 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






      pass







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      pass







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      pass







 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      pass







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      pass







"""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        a     (   unittest.TestCase	):







    def 			UpperCamelCase_      (			self       ):
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()







    def 			UpperCamelCase_      (			self       ):
      lowercase    ,     lowercase							    =       FlaxStableDiffusionPipeline.from_pretrained(
          'stabilityai/stable-diffusion-2'							,		revision='bf16'							,		dtype=jnp.bfloataa							,		)

      lowercase							    =       'A painting of a squirrel eating a burger'
      lowercase							    =       jax.device_count()
      lowercase							    =       num_samples * [prompt]
      lowercase							    =       sd_pipe.prepare_inputs(_lowerCamelCase       )

      lowercase							    =       replicate(_lowerCamelCase       )
      lowercase							    =       shard(_lowerCamelCase       )

      lowercase							    =       jax.random.PRNGKey(0       )
      lowercase							    =       jax.random.split(_lowerCamelCase							,		jax.device_count()       )

      lowercase							    =       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)

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

      lowercase							    =       jnp.asarray(jax.device_get(image_slice.flatten()       )       )
      lowercase							    =       jnp.array([0.4_2_3_8, 0.4_4_1_4, 0.4_3_9_5, 0.4_4_5_3, 0.4_6_2_9, 0.4_5_9_0, 0.4_5_3_1, 0.4_5_5_0_8, 0.4_5_1_2]       )
      print(F'output_slice: {output_slice}'       )
      assert jnp.abs(output_slice - expected_slice       ).max() < 1e-2







    def 			UpperCamelCase_      (			self       ):
      lowercase							    =       'stabilityai/stable-diffusion-2'
      lowercase    ,     lowercase							    =       FlaxDPMSolverMultistepScheduler.from_pretrained(_lowerCamelCase							,		subfolder='scheduler'       )
      lowercase    ,     lowercase							    =       FlaxStableDiffusionPipeline.from_pretrained(
          _lowerCamelCase							,		scheduler=_lowerCamelCase							,		revision='bf16'							,		dtype=jnp.bfloataa							,		)
      lowercase							    =       scheduler_params

      lowercase							    =       'A painting of a squirrel eating a burger'
      lowercase							    =       jax.device_count()
      lowercase							    =       num_samples * [prompt]
      lowercase							    =       sd_pipe.prepare_inputs(_lowerCamelCase       )

      lowercase							    =       replicate(_lowerCamelCase       )
      lowercase							    =       shard(_lowerCamelCase       )

      lowercase							    =       jax.random.PRNGKey(0       )
      lowercase							    =       jax.random.split(_lowerCamelCase							,		jax.device_count()       )

      lowercase							    =       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)

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

      lowercase							    =       jnp.asarray(jax.device_get(image_slice.flatten()       )       )
      lowercase							    =       jnp.array([0.4_3_3_6, 0.4_2_9_6_9, 0.4_4_5_3, 0.4_1_9_9, 0.4_2_9_7, 0.4_5_3_1, 0.4_4_3_4, 0.4_4_3_4, 0.4_2_9_7]       )
      print(F'output_slice: {output_slice}'       )
      assert jnp.abs(output_slice - expected_slice       ).max() < 1e-2


