






"""simple docstring"""
import json
from typing import TYPE_CHECKING, List, Optional, Tuple

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_blenderbot import BlenderbotTokenizer


if TYPE_CHECKING:
			from transformers.pipelines.conversational import Conversation

lowerCamelCase_							      = logging.get_logger(__name__)


lowerCamelCase_							      = {
    '''vocab_file''': '''vocab.json''',
    '''merges_file''': '''merges.txt''',
    '''tokenizer_config_file''': '''tokenizer_config.json''',
}

lowerCamelCase_							      = {
    '''vocab_file''': {'''facebook/blenderbot-3B''': '''https://huggingface.co/facebook/blenderbot-3B/resolve/main/vocab.json'''},
    '''merges_file''': {'''facebook/blenderbot-3B''': '''https://huggingface.co/facebook/blenderbot-3B/resolve/main/merges.txt'''},
    '''tokenizer_config_file''': {
        '''facebook/blenderbot-3B''': '''https://huggingface.co/facebook/blenderbot-3B/resolve/main/tokenizer_config.json'''
    },
}

lowerCamelCase_							      = {'''facebook/blenderbot-3B''': 128}







class      UpperCamelCase_			(__A	):
	__magic_name__       =							VOCAB_FILES_NAMES
	__magic_name__       =							PRETRAINED_VOCAB_FILES_MAP
	__magic_name__       =							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	__magic_name__       =							['''input_ids''', '''attention_mask''']
	__magic_name__       =							BlenderbotTokenizer





	def __init__( self   :					List[str]      ,      lowerCAmelCase_   :					List[Any]=None      ,      lowerCAmelCase_   :					Tuple=None      ,      lowerCAmelCase_   :					str=None      ,      lowerCAmelCase_   :					List[Any]="replace"      ,      lowerCAmelCase_   :					List[Any]="<s>"      ,      lowerCAmelCase_   :					Tuple="</s>"      ,      lowerCAmelCase_   :					str="</s>"      ,      lowerCAmelCase_   :					List[str]="<s>"      ,      lowerCAmelCase_   :					int="<unk>"      ,      lowerCAmelCase_   :					Any="<pad>"      ,      lowerCAmelCase_   :					str="<mask>"      ,      lowerCAmelCase_   :					Optional[int]=False      ,      lowerCAmelCase_   :					Dict=True      ,      **lowerCAmelCase_   :					Any      ,      )		->     Dict:
				super().__init__(
				    lowerCAmelCase_      ,      lowerCAmelCase_      ,      tokenizer_file=lowerCAmelCase_      ,      errors=lowerCAmelCase_      ,      bos_token=lowerCAmelCase_      ,      eos_token=lowerCAmelCase_      ,      sep_token=lowerCAmelCase_      ,      cls_token=lowerCAmelCase_      ,      unk_token=lowerCAmelCase_      ,      pad_token=lowerCAmelCase_      ,      mask_token=lowerCAmelCase_      ,      add_prefix_space=lowerCAmelCase_      ,      trim_offsets=lowerCAmelCase_      ,      **lowerCAmelCase_      ,      )

				UpperCAmelCase_	:					Optional[int]     	=   json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()		)
				if pre_tok_state.get("add_prefix_space"      ,      lowerCAmelCase_		) != add_prefix_space:
							UpperCAmelCase_	:					str     	=   getattr(lowerCAmelCase_      ,      pre_tok_state.pop("type"		)		)
							UpperCAmelCase_	:					Optional[Any]     	=   add_prefix_space
							UpperCAmelCase_	:					Any     	=   pre_tok_class(**lowerCAmelCase_		)

				UpperCAmelCase_	:					Optional[Any]     	=   add_prefix_space

				UpperCAmelCase_	:					Union[str, Any]     	=   "post_processor"
				UpperCAmelCase_	:					Dict     	=   getattr(self.backend_tokenizer      ,      lowerCAmelCase_      ,      lowerCAmelCase_		)
				if tokenizer_component_instance:
							UpperCAmelCase_	:					Any     	=   json.loads(tokenizer_component_instance.__getstate__()		)

							# The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
							if "sep" in state:
										UpperCAmelCase_	:					Optional[Any]     	=   tuple(state["sep"]		)
							if "cls" in state:
										UpperCAmelCase_	:					Union[str, Any]     	=   tuple(state["cls"]		)

							UpperCAmelCase_	:					str     	=   False

							if state.get("add_prefix_space"      ,      lowerCAmelCase_		) != add_prefix_space:
										UpperCAmelCase_	:					List[str]     	=   add_prefix_space
										UpperCAmelCase_	:					Union[str, Any]     	=   True

							if state.get("trim_offsets"      ,      lowerCAmelCase_		) != trim_offsets:
										UpperCAmelCase_	:					Any     	=   trim_offsets
										UpperCAmelCase_	:					Any     	=   True

							if changes_to_apply:
										UpperCAmelCase_	:					int     	=   getattr(lowerCAmelCase_      ,      state.pop("type"		)		)
										UpperCAmelCase_	:					Dict     	=   component_class(**lowerCAmelCase_		)
										setattr(self.backend_tokenizer      ,      lowerCAmelCase_      ,      lowerCAmelCase_		)





	@property
	# Copied from transformers.models.roberta.tokenization_roberta_fast.RobertaTokenizerFast.mask_token with Roberta->Blenderbot, RoBERTa->Blenderbot
	def 	_SCREAMING_SNAKE_CASE ( self   :					List[str]		)		->     str:
				if self._mask_token is None:
							if self.verbose:
										logger.error("Using mask_token, but it is not set yet."		)
							return None
				return str(self._mask_token		)





	@mask_token.setter
	def 	_SCREAMING_SNAKE_CASE ( self   :					Union[str, Any]      ,      lowerCAmelCase_   :					Tuple		)		->     Tuple:
				UpperCAmelCase_	:					Optional[Any]     	=   AddedToken(lowerCAmelCase_      ,      lstrip=lowerCAmelCase_      ,      rstrip=lowerCAmelCase_		) if isinstance(lowerCAmelCase_      ,      lowerCAmelCase_		) else value
				UpperCAmelCase_	:					Any     	=   value





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[int]      ,      *lowerCAmelCase_   :					Optional[Any]      ,      **lowerCAmelCase_   :					List[str]		)		->     BatchEncoding:
				UpperCAmelCase_	:					List[Any]     	=   kwargs.get("is_split_into_words"      ,      lowerCAmelCase_		)
				assert self.add_prefix_space or not is_split_into_words, (
				    f"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True """
				    "to use it with pretokenized inputs."
				)

				return super()._batch_encode_plus(*lowerCAmelCase_      ,      **lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[Any]      ,      *lowerCAmelCase_   :					List[str]      ,      **lowerCAmelCase_   :					Tuple		)		->     BatchEncoding:
				UpperCAmelCase_	:					Optional[int]     	=   kwargs.get("is_split_into_words"      ,      lowerCAmelCase_		)

				assert self.add_prefix_space or not is_split_into_words, (
				    f"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True """
				    "to use it with pretokenized inputs."
				)

				return super()._encode_plus(*lowerCAmelCase_      ,      **lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[Any]      ,      lowerCAmelCase_   :					str      ,      lowerCAmelCase_   :					Optional[str] = None		)		->     Tuple[str]:
				UpperCAmelCase_	:					Union[str, Any]     	=   self._tokenizer.model.save(lowerCAmelCase_      ,      name=lowerCAmelCase_		)
				return tuple(lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Dict      ,      lowerCAmelCase_   :					List[int]      ,      lowerCAmelCase_   :					Optional[List[int]] = None		)		->     List[int]:
				UpperCAmelCase_	:					Dict     	=   [self.sep_token_id]
				UpperCAmelCase_	:					List[str]     	=   [self.cls_token_id]

				if token_ids_a is None:
							return len(cls + token_ids_a + sep		) * [0]
				return len(cls + token_ids_a + sep + sep + token_ids_a + sep		) * [0]





	def 	_SCREAMING_SNAKE_CASE ( self   :					int      ,      lowerCAmelCase_   :					List[int]      ,      lowerCAmelCase_   :					Optional[List[int]] = None		)		->     Optional[Any]:
				return token_ids_a + [self.eos_token_id]





	def 	_SCREAMING_SNAKE_CASE ( self   :					Tuple      ,      lowerCAmelCase_   :					"Conversation"		)		->     List[int]:
				UpperCAmelCase_	:					Tuple     	=   []
				for is_user, text in conversation.iter_texts():
							if is_user:
										# We need to space prefix as it's being done within blenderbot
										inputs.append(" " + text		)
							else:
										# Generated responses should contain them already.
										inputs.append(lowerCAmelCase_		)

				UpperCAmelCase_	:					Dict     	=   "  ".join(lowerCAmelCase_		)
				UpperCAmelCase_	:					List[str]     	=   self.encode(lowerCAmelCase_		)
				if len(lowerCAmelCase_		) > self.model_max_length:
							UpperCAmelCase_	:					Optional[int]     	=   input_ids[-self.model_max_length :]
							logger.warning(f"""Trimmed input from conversation as it was longer than {self.model_max_length} tokens."""		)
				return input_ids








"""simple docstring"""




import os
import tempfile
import unittest

import numpy as np

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


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

 from diffusers import FlaxDDIMScheduler, FlaxDiffusionPipeline, FlaxStableDiffusionPipeline






@require_flax
class       UpperCAmelCase_    (							unittest.TestCase		):



      def 		snake_case_   (	self	:      Optional[int]		):
           with tempfile.TemporaryDirectory() as tmpdirname:
                # pipeline has Flax weights
                _UpperCAmelCase     :      Union[str, Any]														=       FlaxDiffusionPipeline.from_pretrained(
                    "hf-internal-testing/tiny-stable-diffusion-pipe"	,      safety_checker=A	,      cache_dir=A		)

                _UpperCAmelCase     :      List[str]														=       [t[-1] for t in os.walk(os.path.join(A	,      os.listdir(A		)[0]	,      "snapshots"		)		)]
                _UpperCAmelCase     :      Union[str, Any]														=       [item for sublist in all_root_files for item in sublist]

                # None of the downloaded files should be a PyTorch file even if we have some here:
                # https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_pytorch_model.bin
                assert not any(f.endswith(".bin"		) for f in files		)







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



      def 		snake_case_   (	self	:      List[str]		):
           _UpperCAmelCase							,			_UpperCAmelCase     :      int														=       FlaxStableDiffusionPipeline.from_pretrained(
               "hf-internal-testing/tiny-stable-diffusion-pipe"	,      safety_checker=A		)

           _UpperCAmelCase     :      str														=       (
               "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
               " field, close up, split lighting, cinematic"
           )

           _UpperCAmelCase     :      int														=       jax.random.PRNGKey(0		)
           _UpperCAmelCase     :      List[str]														=       4

           _UpperCAmelCase     :      Any														=       jax.device_count()
           _UpperCAmelCase     :      Tuple														=       num_samples * [prompt]
           _UpperCAmelCase     :      Optional[int]														=       pipeline.prepare_inputs(A		)

           # shard inputs and rng
           _UpperCAmelCase     :      Union[str, Any]														=       replicate(A		)
           _UpperCAmelCase     :      List[Any]														=       jax.random.split(A	,      A		)
           _UpperCAmelCase     :      List[Any]														=       shard(A		)

           _UpperCAmelCase     :      Tuple														=       pipeline(A	,      A	,      A	,      A	,      jit=A		).images

           assert images.shape == (num_samples, 1, 6_4, 6_4, 3)
           if jax.device_count() == 8:
                assert np.abs(np.abs(images[0, 0, :2, :2, -2:]	,      dtype=np.floataa		).sum() - 4.1_514_745		) < 1e-3
                assert np.abs(np.abs(A	,      dtype=np.floataa		).sum() - 49_947.875		) < 5e-1

           _UpperCAmelCase     :      Optional[int]														=       pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]		)		)		)
           assert len(A		) == num_samples



      def 		snake_case_   (	self	:      Union[str, Any]		):
           _UpperCAmelCase							,			_UpperCAmelCase     :      Optional[Any]														=       FlaxStableDiffusionPipeline.from_pretrained(
               "CompVis/stable-diffusion-v1-4"	,      revision="flax"	,      safety_checker=A		)

           _UpperCAmelCase     :      Dict														=       (
               "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
               " field, close up, split lighting, cinematic"
           )

           _UpperCAmelCase     :      Dict														=       jax.random.PRNGKey(0		)
           _UpperCAmelCase     :      str														=       5_0

           _UpperCAmelCase     :      List[str]														=       jax.device_count()
           _UpperCAmelCase     :      Any														=       num_samples * [prompt]
           _UpperCAmelCase     :      Dict														=       pipeline.prepare_inputs(A		)

           # shard inputs and rng
           _UpperCAmelCase     :      Dict														=       replicate(A		)
           _UpperCAmelCase     :      str														=       jax.random.split(A	,      A		)
           _UpperCAmelCase     :      Optional[Any]														=       shard(A		)

           _UpperCAmelCase     :      List[Any]														=       pipeline(A	,      A	,      A	,      A	,      jit=A		).images

           assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
           if jax.device_count() == 8:
                assert np.abs((np.abs(images[0, 0, :2, :2, -2:]	,      dtype=np.floataa		).sum() - 0.05_652_401)		) < 1e-3
                assert np.abs((np.abs(A	,      dtype=np.floataa		).sum() - 2_383_808.2)		) < 5e-1



      def 		snake_case_   (	self	:      List[Any]		):
           _UpperCAmelCase							,			_UpperCAmelCase     :      List[Any]														=       FlaxStableDiffusionPipeline.from_pretrained(
               "CompVis/stable-diffusion-v1-4"	,      revision="bf16"	,      dtype=jnp.bfloataa	,      safety_checker=A		)

           _UpperCAmelCase     :      Any														=       (
               "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
               " field, close up, split lighting, cinematic"
           )

           _UpperCAmelCase     :      List[Any]														=       jax.random.PRNGKey(0		)
           _UpperCAmelCase     :      List[Any]														=       5_0

           _UpperCAmelCase     :      Union[str, Any]														=       jax.device_count()
           _UpperCAmelCase     :      Optional[int]														=       num_samples * [prompt]
           _UpperCAmelCase     :      Union[str, Any]														=       pipeline.prepare_inputs(A		)

           # shard inputs and rng
           _UpperCAmelCase     :      int														=       replicate(A		)
           _UpperCAmelCase     :      Optional[int]														=       jax.random.split(A	,      A		)
           _UpperCAmelCase     :      str														=       shard(A		)

           _UpperCAmelCase     :      Union[str, Any]														=       pipeline(A	,      A	,      A	,      A	,      jit=A		).images

           assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
           if jax.device_count() == 8:
                assert np.abs((np.abs(images[0, 0, :2, :2, -2:]	,      dtype=np.floataa		).sum() - 0.04_003_906)		) < 1e-3
                assert np.abs((np.abs(A	,      dtype=np.floataa		).sum() - 2_373_516.75)		) < 5e-1



      def 		snake_case_   (	self	:      int		):
           _UpperCAmelCase							,			_UpperCAmelCase     :      Dict														=       FlaxStableDiffusionPipeline.from_pretrained(
               "CompVis/stable-diffusion-v1-4"	,      revision="bf16"	,      dtype=jnp.bfloataa		)

           _UpperCAmelCase     :      str														=       (
               "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
               " field, close up, split lighting, cinematic"
           )

           _UpperCAmelCase     :      Optional[int]														=       jax.random.PRNGKey(0		)
           _UpperCAmelCase     :      Dict														=       5_0

           _UpperCAmelCase     :      Union[str, Any]														=       jax.device_count()
           _UpperCAmelCase     :      str														=       num_samples * [prompt]
           _UpperCAmelCase     :      int														=       pipeline.prepare_inputs(A		)

           # shard inputs and rng
           _UpperCAmelCase     :      Optional[int]														=       replicate(A		)
           _UpperCAmelCase     :      Tuple														=       jax.random.split(A	,      A		)
           _UpperCAmelCase     :      Dict														=       shard(A		)

           _UpperCAmelCase     :      Optional[Any]														=       pipeline(A	,      A	,      A	,      A	,      jit=A		).images

           assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
           if jax.device_count() == 8:
                assert np.abs((np.abs(images[0, 0, :2, :2, -2:]	,      dtype=np.floataa		).sum() - 0.04_003_906)		) < 1e-3
                assert np.abs((np.abs(A	,      dtype=np.floataa		).sum() - 2_373_516.75)		) < 5e-1



      def 		snake_case_   (	self	:      Optional[Any]		):
           _UpperCAmelCase     :      Dict														=       FlaxDDIMScheduler(
               beta_start=0.00_085	,      beta_end=0.012	,      beta_schedule="scaled_linear"	,      set_alpha_to_one=A	,      steps_offset=1	,      )

           _UpperCAmelCase							,			_UpperCAmelCase     :      Dict														=       FlaxStableDiffusionPipeline.from_pretrained(
               "CompVis/stable-diffusion-v1-4"	,      revision="bf16"	,      dtype=jnp.bfloataa	,      scheduler=A	,      safety_checker=A	,      )
           _UpperCAmelCase     :      Union[str, Any]														=       scheduler.create_state()

           _UpperCAmelCase     :      str														=       scheduler_state

           _UpperCAmelCase     :      Union[str, Any]														=       (
               "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
               " field, close up, split lighting, cinematic"
           )

           _UpperCAmelCase     :      str														=       jax.random.PRNGKey(0		)
           _UpperCAmelCase     :      List[str]														=       5_0

           _UpperCAmelCase     :      Optional[int]														=       jax.device_count()
           _UpperCAmelCase     :      int														=       num_samples * [prompt]
           _UpperCAmelCase     :      Optional[int]														=       pipeline.prepare_inputs(A		)

           # shard inputs and rng
           _UpperCAmelCase     :      List[Any]														=       replicate(A		)
           _UpperCAmelCase     :      Tuple														=       jax.random.split(A	,      A		)
           _UpperCAmelCase     :      Union[str, Any]														=       shard(A		)

           _UpperCAmelCase     :      Dict														=       pipeline(A	,      A	,      A	,      A	,      jit=A		).images

           assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
           if jax.device_count() == 8:
                assert np.abs((np.abs(images[0, 0, :2, :2, -2:]	,      dtype=np.floataa		).sum() - 0.045_043_945)		) < 1e-3
                assert np.abs((np.abs(A	,      dtype=np.floataa		).sum() - 2_347_693.5)		) < 5e-1



      def 		snake_case_   (	self	:      int		):
           _UpperCAmelCase     :      Optional[Any]														=       (
               "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
               " field, close up, split lighting, cinematic"
           )

           _UpperCAmelCase     :      Optional[Any]														=       jax.device_count()
           _UpperCAmelCase     :      Union[str, Any]														=       num_samples * [prompt]
           _UpperCAmelCase     :      Optional[Any]														=       jax.random.split(jax.random.PRNGKey(0		)	,      A		)

           _UpperCAmelCase							,			_UpperCAmelCase     :      str														=       FlaxStableDiffusionPipeline.from_pretrained(
               "CompVis/stable-diffusion-v1-4"	,      revision="bf16"	,      dtype=jnp.bfloataa	,      safety_checker=A	,      )

           _UpperCAmelCase     :      int														=       replicate(A		)
           _UpperCAmelCase     :      Union[str, Any]														=       pipeline.prepare_inputs(A		)
           _UpperCAmelCase     :      int														=       shard(A		)
           _UpperCAmelCase     :      Dict														=       pipeline(A	,      A	,      A	,      jit=A		).images
           assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
           _UpperCAmelCase     :      List[str]														=       images[2, 0, 2_5_6, 1_0:1_7, 1]

           # With memory efficient attention
           _UpperCAmelCase							,			_UpperCAmelCase     :      str														=       FlaxStableDiffusionPipeline.from_pretrained(
               "CompVis/stable-diffusion-v1-4"	,      revision="bf16"	,      dtype=jnp.bfloataa	,      safety_checker=A	,      use_memory_efficient_attention=A	,      )

           _UpperCAmelCase     :      Any														=       replicate(A		)
           _UpperCAmelCase     :      str														=       pipeline.prepare_inputs(A		)
           _UpperCAmelCase     :      Tuple														=       shard(A		)
           _UpperCAmelCase     :      Tuple														=       pipeline(A	,      A	,      A	,      jit=A		).images
           assert images_eff.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
           _UpperCAmelCase     :      str														=       images[2, 0, 2_5_6, 1_0:1_7, 1]

           # I checked the results visually and they are very similar. However, I saw that the max diff is `1` and the `sum`
           # over the 8 images is exactly `256`, which is very suspicious. Testing a random slice for now.
           assert abs(slice_eff - slice		).max() < 1e-2
