





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 							_UpperCamelCase					(	lowerCAmelCase_  ,						unittest.TestCase     ):
       _UpperCamelCase				:					Optional[Any]         =   PriorTransformer
       _UpperCamelCase				:					str         =   '''hidden_states'''

       @property
       def lowercase (		self: Union[str, Any]		)       -> Optional[Any]:




              """simple docstring"""


              UpperCamelCase_          =  4
              UpperCamelCase_          =  8
              UpperCamelCase_          =  7

              UpperCamelCase_          =  floats_tensor((batch_size, embedding_dim)		).to(_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  floats_tensor((batch_size, embedding_dim)		).to(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  floats_tensor((batch_size, num_embeddings, embedding_dim)		).to(_SCREAMING_SNAKE_CASE		)

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

       def lowercase (		self: int					,   _SCREAMING_SNAKE_CASE: Any=0		)       -> Union[str, Any]:




              """simple docstring"""


              torch.manual_seed(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  4
              UpperCamelCase_          =  8
              UpperCamelCase_          =  7

              UpperCamelCase_          =  torch.randn((batch_size, embedding_dim)		).to(_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  torch.randn((batch_size, embedding_dim)		).to(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  torch.randn((batch_size, num_embeddings, embedding_dim)		).to(_SCREAMING_SNAKE_CASE		)

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

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




              """simple docstring"""


              return (4, 8)

       @property
       def lowercase (		self: Tuple		)       -> Optional[int]:




              """simple docstring"""


              return (4, 8)

       def lowercase (		self: int		)       -> Any:




              """simple docstring"""


              UpperCamelCase_          =  {
                  "num_attention_heads": 2,
                  "attention_head_dim": 4,
                  "num_layers": 2,
                  "embedding_dim": 8,
                  "num_embeddings": 7,
                  "additional_embeddings": 4,
              }
              UpperCamelCase_          =  self.dummy_input
              return init_dict, inputs_dict

       def lowercase (		self: str		)       -> Optional[int]:




              """simple docstring"""


              UpperCamelCase_       , UpperCamelCase_          =  PriorTransformer.from_pretrained(
                  "hf-internal-testing/prior-dummy"					,   output_loading_info=_SCREAMING_SNAKE_CASE		)
              self.assertIsNotNone(_SCREAMING_SNAKE_CASE		)
              self.assertEqual(len(loading_info["missing_keys"]		)					,   0		)

              model.to(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  model(**self.dummy_input		)[0]

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

       def lowercase (		self: Tuple		)       -> Dict:




              """simple docstring"""


              UpperCamelCase_       , UpperCamelCase_          =  self.prepare_init_args_and_inputs_for_common()

              UpperCamelCase_          =  self.model_class(**_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  inspect.signature(model.forward		)
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              UpperCamelCase_          =  [*signature.parameters.keys()]

              UpperCamelCase_          =  ["hidden_states", "timestep"]
              self.assertListEqual(arg_names[:2]					,   _SCREAMING_SNAKE_CASE		)



       def lowercase (		self: List[str]		)       -> Dict:




              """simple docstring"""


              UpperCamelCase_          =  PriorTransformer.from_pretrained("hf-internal-testing/prior-dummy"		)
              UpperCamelCase_          =  model.to(_SCREAMING_SNAKE_CASE		)

              if hasattr(_SCREAMING_SNAKE_CASE					,   "set_default_attn_processor"		):
                     model.set_default_attn_processor()

              UpperCamelCase_          =  self.get_dummy_seed_input()

              with torch.no_grad():
                     UpperCamelCase_          =  model(**_SCREAMING_SNAKE_CASE		)[0]

              UpperCamelCase_          =  output[0, :5].flatten().cpu()
              print(_SCREAMING_SNAKE_CASE		)

              # 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.
              UpperCamelCase_          =  torch.tensor([-1.34_36, -0.28_70, 0.75_38, 0.43_68, -0.02_39]		)
              self.assertTrue(torch_all_close(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   rtol=1e-2		)		)



@slow
class 							_UpperCamelCase					(	unittest.TestCase     ):

       def lowercase (		self: Optional[int]					,   _SCREAMING_SNAKE_CASE: str=1					,   _SCREAMING_SNAKE_CASE: List[Any]=768					,   _SCREAMING_SNAKE_CASE: Optional[int]=77					,   _SCREAMING_SNAKE_CASE: Dict=0		)       -> str:




              """simple docstring"""


              torch.manual_seed(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  batch_size
              UpperCamelCase_          =  embedding_dim
              UpperCamelCase_          =  num_embeddings

              UpperCamelCase_          =  torch.randn((batch_size, embedding_dim)		).to(_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  torch.randn((batch_size, embedding_dim)		).to(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  torch.randn((batch_size, num_embeddings, embedding_dim)		).to(_SCREAMING_SNAKE_CASE		)

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

       def lowercase (		self: List[Any]		)       -> List[str]:




              """simple docstring"""


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



       @parameterized.expand(
           [
               # fmt: off
               [13, [-0.58_61,  0.12_83, -0.09_31,  0.08_82,  0.44_76,  0.13_29, -0.04_98,  0.06_40]],
               [37, [-0.49_13,  0.01_10, -0.04_83,  0.05_41,  0.49_54, -0.01_70,  0.03_54,  0.16_51]],
               # fmt: on
           ]		)
       def lowercase (		self: Any					,   _SCREAMING_SNAKE_CASE: Optional[int]					,   _SCREAMING_SNAKE_CASE: List[str]		)       -> Any:




              """simple docstring"""


              UpperCamelCase_          =  PriorTransformer.from_pretrained("kandinsky-community/kandinsky-2-1-prior"					,   subfolder="prior"		)
              model.to(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  self.get_dummy_seed_input(seed=_SCREAMING_SNAKE_CASE		)

              with torch.no_grad():
                     UpperCamelCase_          =  model(**_SCREAMING_SNAKE_CASE		)[0]

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

              UpperCamelCase_          =  sample[0, :8].flatten().cpu()
              print(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  torch.tensor(_SCREAMING_SNAKE_CASE		)

              assert torch_all_close(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   atol=1e-3		)








import re
import string

import numpy as np

import datasets


_UpperCAmelCase        =				'\nReturns the rate at which the input predicted strings exactly match their references, ignoring any strings input as part of the regexes_to_ignore list.\n'

_UpperCAmelCase        =				'\nArgs:\n    predictions: List of predicted texts.\n    references: List of reference texts.\n    regexes_to_ignore: List, defaults to None. Regex expressions of characters to\n        ignore when calculating the exact matches. Note: these regexes are removed\n        from the input data before the changes based on the options below (e.g. ignore_case,\n        ignore_punctuation, ignore_numbers) are applied.\n    ignore_case: Boolean, defaults to False. If true, turns everything\n        to lowercase so that capitalization differences are ignored.\n    ignore_punctuation: Boolean, defaults to False. If true, removes all punctuation before\n        comparing predictions and references.\n    ignore_numbers: Boolean, defaults to False. If true, removes all punctuation before\n        comparing predictions and references.\nReturns:\n    exact_match: Dictionary containing exact_match rate. Possible values are between 0.0 and 100.0, inclusive.\nExamples:\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds)\n    >>> print(round(results["exact_match"], 1))\n    25.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell"], ignore_case=True, ignore_punctuation=True)\n    >>> print(round(results["exact_match"], 1))\n    50.0\n\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell", "YELL"], ignore_case=True, ignore_punctuation=True)\n    >>> print(round(results["exact_match"], 1))\n    75.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell", "YELL"], ignore_case=True, ignore_punctuation=True, ignore_numbers=True)\n    >>> print(round(results["exact_match"], 1))\n    100.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["The cat sat on the mat.", "Theaters are great.", "It\'s like comparing oranges and apples."]\n    >>> preds = ["The cat sat on the mat?", "Theaters are great.", "It\'s like comparing apples and oranges."]\n    >>> results = exact_match.compute(references=refs, predictions=preds)\n    >>> print(round(results["exact_match"], 1))\n    33.3\n\n'

_UpperCAmelCase        =				'\n'



@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION  ,						_KWARGS_DESCRIPTION     )
class 							_UpperCamelCase					(	datasets.Metric     ):

       def lowercase (		self: Any		)       -> List[str]:




              """simple docstring"""


              return datasets.MetricInfo(
                  description=_DESCRIPTION					,   citation=_CITATION					,   inputs_description=_KWARGS_DESCRIPTION					,   features=datasets.Features(
                      {
                          "predictions": datasets.Value("string"					,   id="sequence"		),
                          "references": datasets.Value("string"					,   id="sequence"		),
                      }		)					,   reference_urls=[]					,   )



       def lowercase (		self: Any					,   _SCREAMING_SNAKE_CASE: Any					,   _SCREAMING_SNAKE_CASE: Tuple					,   _SCREAMING_SNAKE_CASE: List[Any]=None					,   _SCREAMING_SNAKE_CASE: Tuple=False					,   _SCREAMING_SNAKE_CASE: Optional[int]=False					,   _SCREAMING_SNAKE_CASE: Any=False					,   )       -> Optional[Any]:




              """simple docstring"""


              if regexes_to_ignore is not None:
                     for s in regexes_to_ignore:
                            UpperCamelCase_          =  np.array([re.sub(_SCREAMING_SNAKE_CASE					,   ""					,   _SCREAMING_SNAKE_CASE		) for x in predictions]		)
                            UpperCamelCase_          =  np.array([re.sub(_SCREAMING_SNAKE_CASE					,   ""					,   _SCREAMING_SNAKE_CASE		) for x in references]		)
              else:
                     UpperCamelCase_          =  np.asarray(_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  np.asarray(_SCREAMING_SNAKE_CASE		)

              if ignore_case:
                     UpperCamelCase_          =  np.char.lower(_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  np.char.lower(_SCREAMING_SNAKE_CASE		)

              if ignore_punctuation:
                     UpperCamelCase_          =  string.punctuation.maketrans(""					,   ""					,   string.punctuation		)
                     UpperCamelCase_          =  np.char.translate(_SCREAMING_SNAKE_CASE					,   table=_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  np.char.translate(_SCREAMING_SNAKE_CASE					,   table=_SCREAMING_SNAKE_CASE		)

              if ignore_numbers:
                     UpperCamelCase_          =  string.digits.maketrans(""					,   ""					,   string.digits		)
                     UpperCamelCase_          =  np.char.translate(_SCREAMING_SNAKE_CASE					,   table=_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  np.char.translate(_SCREAMING_SNAKE_CASE					,   table=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  predictions == references

              return {"exact_match": np.mean(_SCREAMING_SNAKE_CASE		) * 100}

