





import json
import os
from functools import lru_cache
from typing import List, Optional, Tuple

import regex as re

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


_SCREAMING_SNAKE_CASE      	=		logging.get_logger(__name__)


_SCREAMING_SNAKE_CASE      	=		{"""vocab_file""": """vocab.json""", """merges_file""": """merges.txt"""}

# See all BART models at https://huggingface.co/models?filter=bart
_SCREAMING_SNAKE_CASE      	=		{
    """vocab_file""": {
        """facebook/bart-base""": """https://huggingface.co/facebook/bart-base/resolve/main/vocab.json""",
        """facebook/bart-large""": """https://huggingface.co/facebook/bart-large/resolve/main/vocab.json""",
        """facebook/bart-large-mnli""": """https://huggingface.co/facebook/bart-large-mnli/resolve/main/vocab.json""",
        """facebook/bart-large-cnn""": """https://huggingface.co/facebook/bart-large-cnn/resolve/main/vocab.json""",
        """facebook/bart-large-xsum""": """https://huggingface.co/facebook/bart-large-xsum/resolve/main/vocab.json""",
        """yjernite/bart_eli5""": """https://huggingface.co/yjernite/bart_eli5/resolve/main/vocab.json""",
    },
    """merges_file""": {
        """facebook/bart-base""": """https://huggingface.co/facebook/bart-base/resolve/main/merges.txt""",
        """facebook/bart-large""": """https://huggingface.co/facebook/bart-large/resolve/main/merges.txt""",
        """facebook/bart-large-mnli""": """https://huggingface.co/facebook/bart-large-mnli/resolve/main/merges.txt""",
        """facebook/bart-large-cnn""": """https://huggingface.co/facebook/bart-large-cnn/resolve/main/merges.txt""",
        """facebook/bart-large-xsum""": """https://huggingface.co/facebook/bart-large-xsum/resolve/main/merges.txt""",
        """yjernite/bart_eli5""": """https://huggingface.co/yjernite/bart_eli5/resolve/main/merges.txt""",
    },
}

_SCREAMING_SNAKE_CASE      	=		{
    """facebook/bart-base""": 10_24,
    """facebook/bart-large""": 10_24,
    """facebook/bart-large-mnli""": 10_24,
    """facebook/bart-large-cnn""": 10_24,
    """facebook/bart-large-xsum""": 10_24,
    """yjernite/bart_eli5""": 10_24,
}


@lru_cache()
def 							SCREAMING_SNAKE_CASE__						(    ):
 snake_case_		:			Optional[Any]											=	(
     list(range(ord('!'   )   , ord('~'   ) + 1   )   ) + list(range(ord('¡'   )   , ord('¬'   ) + 1   )   ) + list(range(ord('®'   )   , ord('ÿ'   ) + 1   )   )
 )
 snake_case_		:			List[Any]											=	bs[:]
 snake_case_		:			Dict											=	0
 for b in range(2**8   ):
  if b not in bs:
   bs.append(__a   )
   cs.append(2**8 + n   )
   n += 1
 snake_case_		:			str											=	[chr(__a   ) for n in cs]
 return dict(zip(__a   , __a   )   )


def 							SCREAMING_SNAKE_CASE__						(    __a   ):
 snake_case_		:			List[Any]											=	set()
 snake_case_		:			List[str]											=	word[0]
 for char in word[1:]:
  pairs.add((prev_char, char)   )
  snake_case_		:			Any											=	char
 return pairs


class SCREAMING_SNAKE_CASE_						(					snake_case_					):
      __magic_name__:			Any       =       VOCAB_FILES_NAMES
      __magic_name__:			Tuple       =       PRETRAINED_VOCAB_FILES_MAP
      __magic_name__:			str       =       PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __magic_name__:			List[Any]       =       ["input_ids", "attention_mask"]






      def __init__(      self      :      Dict		,   _A      :      List[Any]		,   _A      :      Optional[int]		,   _A      :      List[Any]="replace"		,   _A      :      List[str]="<s>"		,   _A      :      Tuple="</s>"		,   _A      :      Optional[int]="</s>"		,   _A      :      Optional[Any]="<s>"		,   _A      :      Tuple="<unk>"		,   _A      :      Any="<pad>"		,   _A      :      List[Any]="<mask>"		,   _A      :      int=False		,   **_A      :      int		,   )							->    int:







       """simple docstring"""


       snake_case_		:			Any											=	AddedToken(_A		,   lstrip=_A		,   rstrip=_A	) if isinstance(_A		,   _A	) else bos_token
       snake_case_		:			Any											=	AddedToken(_A		,   lstrip=_A		,   rstrip=_A	) if isinstance(_A		,   _A	) else eos_token
       snake_case_		:			Dict											=	AddedToken(_A		,   lstrip=_A		,   rstrip=_A	) if isinstance(_A		,   _A	) else sep_token
       snake_case_		:			str											=	AddedToken(_A		,   lstrip=_A		,   rstrip=_A	) if isinstance(_A		,   _A	) else cls_token
       snake_case_		:			Dict											=	AddedToken(_A		,   lstrip=_A		,   rstrip=_A	) if isinstance(_A		,   _A	) else unk_token
       snake_case_		:			Tuple											=	AddedToken(_A		,   lstrip=_A		,   rstrip=_A	) if isinstance(_A		,   _A	) else pad_token

       # Mask token behave like a normal word, i.e. include the space before it
       snake_case_		:			Optional[Any]											=	AddedToken(_A		,   lstrip=_A		,   rstrip=_A	) if isinstance(_A		,   _A	) else mask_token

       super().__init__(
           errors=_A		,   bos_token=_A		,   eos_token=_A		,   unk_token=_A		,   sep_token=_A		,   cls_token=_A		,   pad_token=_A		,   mask_token=_A		,   add_prefix_space=_A		,   **_A		,   )

       with open(_A		,   encoding='utf-8'	) as vocab_handle:
        snake_case_		:			Optional[Any]											=	json.load(_A	)
       snake_case_		:			str											=	{v: k for k, v in self.encoder.items()}
       snake_case_		:			str											=	errors  # how to handle errors in decoding
       snake_case_		:			str											=	bytes_to_unicode()
       snake_case_		:			Tuple											=	{v: k for k, v in self.byte_encoder.items()}
       with open(_A		,   encoding='utf-8'	) as merges_handle:
        snake_case_		:			Optional[int]											=	merges_handle.read().split('\n'	)[1:-1]
       snake_case_		:			Optional[Any]											=	[tuple(merge.split()	) for merge in bpe_merges]
       snake_case_		:			List[Any]											=	dict(zip(_A		,   range(len(_A	)	)	)	)
       snake_case_		:			str											=	{}
       snake_case_		:			List[Any]											=	add_prefix_space

       # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
       snake_case_		:			List[Any]											=	re.compile(R'\'s|\'t|\'re|\'ve|\'m|\'ll|\'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+'	)






      @property
      def UpperCAmelCase_				(      self      :      Optional[Any]	)							->    str:







       """simple docstring"""


       return len(self.encoder	)






      def UpperCAmelCase_				(      self      :      Dict	)							->    str:







       """simple docstring"""


       return dict(self.encoder		,   **self.added_tokens_encoder	)






      def UpperCAmelCase_				(      self      :      Union[str, Any]		,   _A      :      Optional[int]	)							->    Any:







       """simple docstring"""


       if token in self.cache:
        return self.cache[token]
       snake_case_		:			Any											=	tuple(_A	)
       snake_case_		:			str											=	get_pairs(_A	)

       if not pairs:
        return token

       while True:
        snake_case_		:			Tuple											=	min(_A		,   key=lambda _A	: self.bpe_ranks.get(_A		,   float('inf'	)	)	)
        if bigram not in self.bpe_ranks:
         break
        snake_case_		:			Optional[int]											=	bigram
        snake_case_		:			List[str]											=	[]
        snake_case_		:			Optional[int]											=	0
        while i < len(_A	):
         try:
          snake_case_		:			Tuple											=	word.index(_A		,   _A	)
         except ValueError:
          new_word.extend(word[i:]	)
          break
         else:
          new_word.extend(word[i:j]	)
          snake_case_		:			int											=	j

         if word[i] == first and i < len(_A	) - 1 and word[i + 1] == second:
          new_word.append(first + second	)
          i += 2
         else:
          new_word.append(word[i]	)
          i += 1
        snake_case_		:			str											=	tuple(_A	)
        snake_case_		:			Optional[int]											=	new_word
        if len(_A	) == 1:
         break
        else:
         snake_case_		:			Any											=	get_pairs(_A	)
       snake_case_		:			Dict											=	' '.join(_A	)
       snake_case_		:			Dict											=	word
       return word






      def UpperCAmelCase_				(      self      :      List[Any]		,   _A      :      Dict	)							->    Tuple:







       """simple docstring"""


       snake_case_		:			Union[str, Any]											=	[]
       for token in re.findall(self.pat		,   _A	):
        snake_case_		:			Union[str, Any]											=	''.join(
            self.byte_encoder[b] for b in token.encode('utf-8'	)	)  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
        bpe_tokens.extend(bpe_token for bpe_token in self.bpe(_A	).split(' '	)	)
       return bpe_tokens






      def UpperCAmelCase_				(      self      :      List[str]		,   _A      :      Optional[int]	)							->    Union[str, Any]:







       """simple docstring"""


       return self.encoder.get(_A		,   self.encoder.get(self.unk_token	)	)






      def UpperCAmelCase_				(      self      :      Any		,   _A      :      str	)							->    Union[str, Any]:







       """simple docstring"""


       return self.decoder.get(_A	)






      def UpperCAmelCase_				(      self      :      str		,   _A      :      int	)							->    Dict:







       """simple docstring"""


       snake_case_		:			Union[str, Any]											=	''.join(_A	)
       snake_case_		:			Any											=	bytearray([self.byte_decoder[c] for c in text]	).decode('utf-8'		,   errors=self.errors	)
       return text






      def UpperCAmelCase_				(      self      :      Dict		,   _A      :      str		,   _A      :      Optional[str] = None	)							->    Tuple[str]:







       """simple docstring"""


       if not os.path.isdir(_A	):
        logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""	)
        return
       snake_case_		:			List[Any]											=	os.path.join(
           _A		,   (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']	)
       snake_case_		:			Tuple											=	os.path.join(
           _A		,   (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['merges_file']	)

       with open(_A		,   'w'		,   encoding='utf-8'	) as f:
        f.write(json.dumps(self.encoder		,   indent=2		,   sort_keys=_A		,   ensure_ascii=_A	) + '\n'	)

       snake_case_		:			Dict											=	0
       with open(_A		,   'w'		,   encoding='utf-8'	) as writer:
        writer.write('#version: 0.2\n'	)
        for bpe_tokens, token_index in sorted(self.bpe_ranks.items()		,   key=lambda _A	: kv[1]	):
         if index != token_index:
          logger.warning(
              F"""Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."""
              ' Please check that the tokenizer is not corrupted!'	)
          snake_case_		:			Optional[Any]											=	token_index
         writer.write(' '.join(_A	) + '\n'	)
         index += 1

       return vocab_file, merge_file






      def UpperCAmelCase_				(      self      :      Tuple		,   _A      :      List[int]		,   _A      :      Optional[List[int]] = None	)							->    List[int]:







       """simple docstring"""


       if token_ids_a is None:
        return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
       snake_case_		:			Any											=	[self.cls_token_id]
       snake_case_		:			List[Any]											=	[self.sep_token_id]
       return cls + token_ids_a + sep + sep + token_ids_a + sep






      def UpperCAmelCase_				(      self      :      List[Any]		,   _A      :      List[int]		,   _A      :      Optional[List[int]] = None		,   _A      :      bool = False	)							->    List[int]:







       """simple docstring"""


       if already_has_special_tokens:
        return super().get_special_tokens_mask(
            token_ids_a=_A		,   token_ids_a=_A		,   already_has_special_tokens=_A	)

       if token_ids_a is None:
        return [1] + ([0] * len(_A	)) + [1]
       return [1] + ([0] * len(_A	)) + [1, 1] + ([0] * len(_A	)) + [1]






      def UpperCAmelCase_				(      self      :      Union[str, Any]		,   _A      :      List[int]		,   _A      :      Optional[List[int]] = None	)							->    List[int]:







       """simple docstring"""


       snake_case_		:			Union[str, Any]											=	[self.sep_token_id]
       snake_case_		:			Optional[Any]											=	[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 UpperCAmelCase_				(      self      :      List[str]		,   _A      :      Optional[int]		,   _A      :      Dict=False		,   **_A      :      Optional[Any]	)							->    Tuple:







       """simple docstring"""


       snake_case_		:			int											=	kwargs.pop('add_prefix_space'		,   self.add_prefix_space	)
       if (is_split_into_words or add_prefix_space) and (len(_A	) > 0 and not text[0].isspace()):
        snake_case_		:			Optional[int]											=	' ' + text
       return (text, kwargs)








import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()
class SCREAMING_SNAKE_CASE_						(					snake_case_       ,				unittest.TestCase					):
						__magic_name__:			Dict       =       PriorTransformer
						__magic_name__:			str       =       "hidden_states"






						@property
						def UpperCAmelCase_				(      self      :      Optional[Any]	)							->    Optional[int]:







							"""simple docstring"""


							snake_case_		:			Any											=	4
							snake_case_		:			int											=	8
							snake_case_		:			Dict											=	7

							snake_case_		:			Union[str, Any]											=	floats_tensor((batch_size, embedding_dim)	).to(_A	)

							snake_case_		:			int											=	floats_tensor((batch_size, embedding_dim)	).to(_A	)
							snake_case_		:			str											=	floats_tensor((batch_size, num_embeddings, embedding_dim)	).to(_A	)

							return {
							    "hidden_states": hidden_states,
							    "timestep": 2,
							    "proj_embedding": proj_embedding,
							    "encoder_hidden_states": encoder_hidden_states,
							}






						def UpperCAmelCase_				(      self      :      List[Any]		,   _A      :      List[Any]=0	)							->    str:







							"""simple docstring"""


							torch.manual_seed(_A	)
							snake_case_		:			List[Any]											=	4
							snake_case_		:			str											=	8
							snake_case_		:			Any											=	7

							snake_case_		:			List[Any]											=	torch.randn((batch_size, embedding_dim)	).to(_A	)

							snake_case_		:			Optional[int]											=	torch.randn((batch_size, embedding_dim)	).to(_A	)
							snake_case_		:			List[Any]											=	torch.randn((batch_size, num_embeddings, embedding_dim)	).to(_A	)

							return {
							    "hidden_states": hidden_states,
							    "timestep": 2,
							    "proj_embedding": proj_embedding,
							    "encoder_hidden_states": encoder_hidden_states,
							}






						@property
						def UpperCAmelCase_				(      self      :      str	)							->    Optional[Any]:







							"""simple docstring"""


							return (4, 8)






						@property
						def UpperCAmelCase_				(      self      :      Optional[int]	)							->    List[Any]:







							"""simple docstring"""


							return (4, 8)






						def UpperCAmelCase_				(      self      :      Optional[int]	)							->    Union[str, Any]:







							"""simple docstring"""


							snake_case_		:			Union[str, Any]											=	{
							    'num_attention_heads': 2,
							    'attention_head_dim': 4,
							    'num_layers': 2,
							    'embedding_dim': 8,
							    'num_embeddings': 7,
							    'additional_embeddings': 4,
							}
							snake_case_		:			Tuple											=	self.dummy_input
							return init_dict, inputs_dict






						def UpperCAmelCase_				(      self      :      List[Any]	)							->    Dict:







							"""simple docstring"""


							snake_case_						,snake_case_		:			str											=	PriorTransformer.from_pretrained(
							    'hf-internal-testing/prior-dummy'		,   output_loading_info=_A	)
							self.assertIsNotNone(_A	)
							self.assertEqual(len(loading_info['missing_keys']	)		,   0	)

							model.to(_A	)
							snake_case_		:			Optional[Any]											=	model(**self.dummy_input	)[0]

							assert hidden_states is not None, "Make sure output is not None"






						def UpperCAmelCase_				(      self      :      Tuple	)							->    List[Any]:







							"""simple docstring"""


							snake_case_						,snake_case_		:			Optional[int]											=	self.prepare_init_args_and_inputs_for_common()

							snake_case_		:			Tuple											=	self.model_class(**_A	)
							snake_case_		:			Any											=	inspect.signature(model.forward	)
							# signature.parameters is an OrderedDict => so arg_names order is deterministic
							snake_case_		:			List[str]											=	[*signature.parameters.keys()]

							snake_case_		:			int											=	['hidden_states', 'timestep']
							self.assertListEqual(arg_names[:2]		,   _A	)






						def UpperCAmelCase_				(      self      :      Union[str, Any]	)							->    List[Any]:







							"""simple docstring"""


							snake_case_		:			str											=	PriorTransformer.from_pretrained('hf-internal-testing/prior-dummy'	)
							snake_case_		:			str											=	model.to(_A	)

							if hasattr(_A		,   'set_default_attn_processor'	):
								model.set_default_attn_processor()

							snake_case_		:			Optional[int]											=	self.get_dummy_seed_input()

							with torch.no_grad():
								snake_case_		:			Any											=	model(**_A	)[0]

							snake_case_		:			Any											=	output[0, :5].flatten().cpu()
							print(_A	)

							# Since the VAE Gaussian prior's generator is seeded on the appropriate device,
							# the expected output slices are not the same for CPU and GPU.
							snake_case_		:			str											=	torch.tensor([-1.3_4_3_6, -0.2_8_7_0, 0.7_5_3_8, 0.4_3_6_8, -0.0_2_3_9]	)
							self.assertTrue(torch_all_close(_A		,   _A		,   rtol=1E-2	)	)
@slow
class SCREAMING_SNAKE_CASE_						(					unittest.TestCase					):






						def UpperCAmelCase_				(      self      :      Optional[Any]		,   _A      :      int=1		,   _A      :      int=768		,   _A      :      str=77		,   _A      :      List[str]=0	)							->    Tuple:







							"""simple docstring"""


							torch.manual_seed(_A	)
							snake_case_		:			Dict											=	batch_size
							snake_case_		:			Any											=	embedding_dim
							snake_case_		:			int											=	num_embeddings

							snake_case_		:			Dict											=	torch.randn((batch_size, embedding_dim)	).to(_A	)

							snake_case_		:			List[Any]											=	torch.randn((batch_size, embedding_dim)	).to(_A	)
							snake_case_		:			Optional[int]											=	torch.randn((batch_size, num_embeddings, embedding_dim)	).to(_A	)

							return {
							    "hidden_states": hidden_states,
							    "timestep": 2,
							    "proj_embedding": proj_embedding,
							    "encoder_hidden_states": encoder_hidden_states,
							}






						def UpperCAmelCase_				(      self      :      List[Any]	)							->    Optional[int]:







							"""simple docstring"""


							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()






						@parameterized.expand(
						    [
						        # fmt: off
						        [13, [-0.5_8_6_1,  0.1_2_8_3, -0.0_9_3_1,  0.0_8_8_2,  0.4_4_7_6,  0.1_3_2_9, -0.0_4_9_8,  0.0_6_4_0]],
						        [37, [-0.4_9_1_3,  0.0_1_1_0, -0.0_4_8_3,  0.0_5_4_1,  0.4_9_5_4, -0.0_1_7_0,  0.0_3_5_4,  0.1_6_5_1]],
						        # fmt: on
						    ]	)
						def UpperCAmelCase_				(      self      :      Tuple		,   _A      :      List[Any]		,   _A      :      List[str]	)							->    Optional[int]:







							"""simple docstring"""


							snake_case_		:			str											=	PriorTransformer.from_pretrained('kandinsky-community/kandinsky-2-1-prior'		,   subfolder='prior'	)
							model.to(_A	)
							snake_case_		:			Optional[Any]											=	self.get_dummy_seed_input(seed=_A	)

							with torch.no_grad():
								snake_case_		:			str											=	model(**_A	)[0]

							assert list(sample.shape	) == [1, 768]

							snake_case_		:			Optional[Any]											=	sample[0, :8].flatten().cpu()
							print(_A	)
							snake_case_		:			int											=	torch.tensor(_A	)

							assert torch_all_close(_A		,   _A		,   atol=1E-3	)

