





from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..utils import maybe_allow_in_graph
from .activations import get_activation
from .attention_processor import Attention
from .embeddings import CombinedTimestepLabelEmbeddings

@maybe_allow_in_graph
class 		SCREAMING_SNAKE_CASE__    (	nn.Module  ):


      def __init__(self	:				Tuple    ,      a__	:				int    ,      a__	:				int    ,      a__	:				int    ,      a__	:				List[Any]=0.0    ,      a__	:				Optional[int] = None    ,      a__	:				str = "geglu"    ,      a__	:				Optional[int] = None    ,      a__	:				bool = False    ,      a__	:				bool = False    ,      a__	:				bool = False    ,      a__	:				bool = False    ,      a__	:				bool = True    ,      a__	:				str = "layer_norm"    ,      a__	:				bool = False    ,      ):





         """simple docstring"""



         super().__init__()
         __snake_case								=       only_cross_attention

         __snake_case								=       (num_embeds_ada_norm is not None) and norm_type == '''ada_norm_zero'''
         __snake_case								=       (num_embeds_ada_norm is not None) and norm_type == '''ada_norm'''

         if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
            raise ValueError(
                f"""`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"""
                f""" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."""				)

         # Define 3 blocks. Each block has its own normalization layer.
         # 1. Self-Attn
         if self.use_ada_layer_norm:
            __snake_case								=       AdaLayerNorm(a__    ,      a__				)
         elif self.use_ada_layer_norm_zero:
            __snake_case								=       AdaLayerNormZero(a__    ,      a__				)
         else:
            __snake_case								=       nn.LayerNorm(a__    ,      elementwise_affine=a__				)
         __snake_case								=       Attention(
             query_dim=a__    ,      heads=a__    ,      dim_head=a__    ,      dropout=a__    ,      bias=a__    ,      cross_attention_dim=cross_attention_dim if only_cross_attention else None    ,      upcast_attention=a__    ,      )

         # 2. Cross-Attn
         if cross_attention_dim is not None or double_self_attention:
            # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
            # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
            # the second cross attention block.
            __snake_case								=       (
                AdaLayerNorm(a__    ,      a__				)
                if self.use_ada_layer_norm
                else nn.LayerNorm(a__    ,      elementwise_affine=a__				)
            )
            __snake_case								=       Attention(
                query_dim=a__    ,      cross_attention_dim=cross_attention_dim if not double_self_attention else None    ,      heads=a__    ,      dim_head=a__    ,      dropout=a__    ,      bias=a__    ,      upcast_attention=a__    ,      )  # is self-attn if encoder_hidden_states is none
         else:
            __snake_case								=       None
            __snake_case								=       None

         # 3. Feed-forward
         __snake_case								=       nn.LayerNorm(a__    ,      elementwise_affine=a__				)
         __snake_case								=       FeedForward(a__    ,      dropout=a__    ,      activation_fn=a__    ,      final_dropout=a__				)

         # let chunk size default to None
         __snake_case								=       None
         __snake_case								=       0


      def        a			(self	:				Tuple    ,      a__	:				Optional[int]    ,      a__	:				int				):





         """simple docstring"""



         __snake_case								=       chunk_size
         __snake_case								=       dim


      def        a			(self	:				str    ,      a__	:				torch.FloatTensor    ,      a__	:				Optional[torch.FloatTensor] = None    ,      a__	:				Optional[torch.FloatTensor] = None    ,      a__	:				Optional[torch.FloatTensor] = None    ,      a__	:				Optional[torch.LongTensor] = None    ,      a__	:				Dict[str, Any] = None    ,      a__	:				Optional[torch.LongTensor] = None    ,      ):





         """simple docstring"""



         if self.use_ada_layer_norm:
            __snake_case								=       self.norma(a__    ,      a__				)
         elif self.use_ada_layer_norm_zero:
            __snake_case	,     __snake_case	,     __snake_case	,     __snake_case	,     __snake_case								=       self.norma(
                a__    ,      a__    ,      a__    ,      hidden_dtype=hidden_states.dtype				)
         else:
            __snake_case								=       self.norma(a__				)

         __snake_case								=       cross_attention_kwargs if cross_attention_kwargs is not None else {}

         __snake_case								=       self.attna(
             a__    ,      encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None    ,      attention_mask=a__    ,      **a__    ,      )
         if self.use_ada_layer_norm_zero:
            __snake_case								=       gate_msa.unsqueeze(1				) * attn_output
         __snake_case								=       attn_output + hidden_states

         # 2. Cross-Attention
         if self.attna is not None:
            __snake_case								=       (
                self.norma(a__    ,      a__				) if self.use_ada_layer_norm else self.norma(a__				)
            )

            __snake_case								=       self.attna(
                a__    ,      encoder_hidden_states=a__    ,      attention_mask=a__    ,      **a__    ,      )
            __snake_case								=       attn_output + hidden_states

         # 3. Feed-forward
         __snake_case								=       self.norma(a__				)

         if self.use_ada_layer_norm_zero:
            __snake_case								=       norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

         if self._chunk_size is not None:
            # "feed_forward_chunk_size" can be used to save memory
            if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
               raise ValueError(
                   f"""`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."""				)

            __snake_case								=       norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
            __snake_case								=       torch.cat(
                [self.ff(a__				) for hid_slice in norm_hidden_states.chunk(a__    ,      dim=self._chunk_dim				)]    ,      dim=self._chunk_dim    ,      )
         else:
            __snake_case								=       self.ff(a__				)

         if self.use_ada_layer_norm_zero:
            __snake_case								=       gate_mlp.unsqueeze(1				) * ff_output

         __snake_case								=       ff_output + hidden_states

         return hidden_states






class 		SCREAMING_SNAKE_CASE__    (	nn.Module  ):


      def __init__(self	:				Any    ,      a__	:				int    ,      a__	:				Optional[int] = None    ,      a__	:				int = 4    ,      a__	:				float = 0.0    ,      a__	:				str = "geglu"    ,      a__	:				bool = False    ,      ):





         """simple docstring"""



         super().__init__()
         __snake_case								=       int(dim * mult				)
         __snake_case								=       dim_out if dim_out is not None else dim

         if activation_fn == "gelu":
            __snake_case								=       GELU(a__    ,      a__				)
         if activation_fn == "gelu-approximate":
            __snake_case								=       GELU(a__    ,      a__    ,      approximate='''tanh'''				)
         elif activation_fn == "geglu":
            __snake_case								=       GEGLU(a__    ,      a__				)
         elif activation_fn == "geglu-approximate":
            __snake_case								=       ApproximateGELU(a__    ,      a__				)

         __snake_case								=       nn.ModuleList([]				)
         # project in
         self.net.append(a__				)
         # project dropout
         self.net.append(nn.Dropout(a__				)				)
         # project out
         self.net.append(nn.Linear(a__    ,      a__				)				)
         # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
         if final_dropout:
            self.net.append(nn.Dropout(a__				)				)


      def        a			(self	:				str    ,      a__	:				int				):





         """simple docstring"""



         for module in self.net:
            __snake_case								=       module(a__				)
         return hidden_states






class 		SCREAMING_SNAKE_CASE__    (	nn.Module  ):


      def __init__(self	:				int    ,      a__	:				int    ,      a__	:				int    ,      a__	:				str = "none"				):





         """simple docstring"""



         super().__init__()
         __snake_case								=       nn.Linear(a__    ,      a__				)
         __snake_case								=       approximate


      def        a			(self	:				Optional[int]    ,      a__	:				Tuple				):





         """simple docstring"""



         if gate.device.type != "mps":
            return F.gelu(a__    ,      approximate=self.approximate				)
         # mps: gelu is not implemented for float16
         return F.gelu(gate.to(dtype=torch.floataa				)    ,      approximate=self.approximate				).to(dtype=gate.dtype				)


      def        a			(self	:				Any    ,      a__	:				Optional[Any]				):





         """simple docstring"""



         __snake_case								=       self.proj(a__				)
         __snake_case								=       self.gelu(a__				)
         return hidden_states






class 		SCREAMING_SNAKE_CASE__    (	nn.Module  ):


      def __init__(self	:				Dict    ,      a__	:				int    ,      a__	:				int				):





         """simple docstring"""



         super().__init__()
         __snake_case								=       nn.Linear(a__    ,      dim_out * 2				)


      def        a			(self	:				Any    ,      a__	:				List[Any]				):





         """simple docstring"""



         if gate.device.type != "mps":
            return F.gelu(a__				)
         # mps: gelu is not implemented for float16
         return F.gelu(gate.to(dtype=torch.floataa				)				).to(dtype=gate.dtype				)


      def        a			(self	:				str    ,      a__	:				Dict				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       self.proj(a__				).chunk(2    ,      dim=-1				)
         return hidden_states * self.gelu(a__				)






class 		SCREAMING_SNAKE_CASE__    (	nn.Module  ):


      def __init__(self	:				str    ,      a__	:				int    ,      a__	:				int				):





         """simple docstring"""



         super().__init__()
         __snake_case								=       nn.Linear(a__    ,      a__				)


      def        a			(self	:				Tuple    ,      a__	:				Tuple				):





         """simple docstring"""



         __snake_case								=       self.proj(a__				)
         return x * torch.sigmoid(1.7_0_2 * x				)






class 		SCREAMING_SNAKE_CASE__    (	nn.Module  ):


      def __init__(self	:				Tuple    ,      a__	:				Tuple    ,      a__	:				Optional[int]				):





         """simple docstring"""



         super().__init__()
         __snake_case								=       nn.Embedding(a__    ,      a__				)
         __snake_case								=       nn.SiLU()
         __snake_case								=       nn.Linear(a__    ,      embedding_dim * 2				)
         __snake_case								=       nn.LayerNorm(a__    ,      elementwise_affine=a__				)


      def        a			(self	:				Optional[Any]    ,      a__	:				Tuple    ,      a__	:				Union[str, Any]				):





         """simple docstring"""



         __snake_case								=       self.linear(self.silu(self.emb(a__				)				)				)
         __snake_case	,     __snake_case								=       torch.chunk(a__    ,      2				)
         __snake_case								=       self.norm(a__				) * (1 + scale) + shift
         return x






class 		SCREAMING_SNAKE_CASE__    (	nn.Module  ):


      def __init__(self	:				List[Any]    ,      a__	:				Tuple    ,      a__	:				List[str]				):





         """simple docstring"""



         super().__init__()

         __snake_case								=       CombinedTimestepLabelEmbeddings(a__    ,      a__				)

         __snake_case								=       nn.SiLU()
         __snake_case								=       nn.Linear(a__    ,      6 * embedding_dim    ,      bias=a__				)
         __snake_case								=       nn.LayerNorm(a__    ,      elementwise_affine=a__    ,      eps=1E-6				)


      def        a			(self	:				Optional[Any]    ,      a__	:				Union[str, Any]    ,      a__	:				List[str]    ,      a__	:				Dict    ,      a__	:				List[str]=None				):





         """simple docstring"""



         __snake_case								=       self.linear(self.silu(self.emb(a__    ,      a__    ,      hidden_dtype=a__				)				)				)
         __snake_case	,     __snake_case	,     __snake_case	,     __snake_case	,     __snake_case	,     __snake_case								=       emb.chunk(6    ,      dim=1				)
         __snake_case								=       self.norm(a__				) * (1 + scale_msa[:, None]) + shift_msa[:, None]
         return x, gate_msa, shift_mlp, scale_mlp, gate_mlp






class 		SCREAMING_SNAKE_CASE__    (	nn.Module  ):


      def __init__(self	:				Optional[int]    ,      a__	:				int    ,      a__	:				int    ,      a__	:				int    ,      a__	:				Optional[str] = None    ,      a__	:				float = 1E-5				):





         """simple docstring"""



         super().__init__()
         __snake_case								=       num_groups
         __snake_case								=       eps

         if act_fn is None:
            __snake_case								=       None
         else:
            __snake_case								=       get_activation(a__				)

         __snake_case								=       nn.Linear(a__    ,      out_dim * 2				)


      def        a			(self	:				List[str]    ,      a__	:				str    ,      a__	:				Dict				):





         """simple docstring"""



         if self.act:
            __snake_case								=       self.act(a__				)
         __snake_case								=       self.linear(a__				)
         __snake_case								=       emb[:, :, None, None]
         __snake_case	,     __snake_case								=       emb.chunk(2    ,      dim=1				)

         __snake_case								=       F.group_norm(a__    ,      self.num_groups    ,      eps=self.eps				)
         __snake_case								=       x * (1 + scale) + shift
         return x








import re
import string

import numpy as np

import datasets


snake_case_           =      '\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'

snake_case_           =      '\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'

snake_case_           =      '\n'

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


      def        a			(self	:				Optional[Any]				):





         """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        a			(self	:				List[Any]    ,      a__	:				int    ,      a__	:				Optional[int]    ,      a__	:				Optional[Any]=None    ,      a__	:				Any=False    ,      a__	:				Dict=False    ,      a__	:				Tuple=False    ,      ):





         """simple docstring"""



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

         if ignore_case:
            __snake_case								=       np.char.lower(a__				)
            __snake_case								=       np.char.lower(a__				)

         if ignore_punctuation:
            __snake_case								=       string.punctuation.maketrans(''''''    ,      ''''''    ,      string.punctuation				)
            __snake_case								=       np.char.translate(a__    ,      table=a__				)
            __snake_case								=       np.char.translate(a__    ,      table=a__				)

         if ignore_numbers:
            __snake_case								=       string.digits.maketrans(''''''    ,      ''''''    ,      string.digits				)
            __snake_case								=       np.char.translate(a__    ,      table=a__				)
            __snake_case								=       np.char.translate(a__    ,      table=a__				)

         __snake_case								=       predictions == references

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

