




from dataclasses import dataclass
from typing import Dict, Optional, Union

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

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin




@dataclass
class 		__A   (      UpperCamelCase__      ):
						UpperCamelCase      =       42




class 		__A   (      UpperCamelCase__       ,    UpperCamelCase__      ):

						@register_to_config
						def __init__(		self :Optional[int]     ,			__snake_case :int = 32     ,			__snake_case :int = 64     ,			__snake_case :int = 20     ,			__snake_case :int = 7_68     ,			__snake_case :List[Any]=77     ,			__snake_case :Union[str, Any]=4     ,			__snake_case :float = 0.0     ,			__snake_case :str = "silu"     ,			__snake_case :Optional[str] = None     ,			__snake_case :Optional[str] = None     ,			__snake_case :Optional[str] = "linear"     ,			__snake_case :Optional[str] = "prd"     ,			__snake_case :Optional[int] = None     ,			__snake_case :Optional[int] = None     ,			__snake_case :Optional[int] = None     ,			):
										'''simple docstring'''

										super().__init__()
										__magic_name__						:						Tuple										=num_attention_heads
										__magic_name__						:						Tuple										=attention_head_dim
										__magic_name__						:						Optional[int]										=num_attention_heads * attention_head_dim
										__magic_name__						:						Union[str, Any]										=additional_embeddings

										__magic_name__						:						Tuple										=time_embed_dim or inner_dim
										__magic_name__						:						str										=embedding_proj_dim or embedding_dim
										__magic_name__						:						List[Any]										=clip_embed_dim or embedding_dim

										__magic_name__						:						Union[str, Any]										=Timesteps(__snake_case     ,			__snake_case     ,			0			)
										__magic_name__						:						Dict										=TimestepEmbedding(__snake_case     ,			__snake_case     ,			out_dim=__snake_case     ,			act_fn=__snake_case			)

										__magic_name__						:						Tuple										=nn.Linear(__snake_case     ,			__snake_case			)

										if embedding_proj_norm_type is None:
														__magic_name__						:						Tuple										=None
										elif embedding_proj_norm_type == "layer":
														__magic_name__						:						Tuple										=nn.LayerNorm(__snake_case			)
										else:
														raise ValueError(f"unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"			)

										__magic_name__						:						Tuple										=nn.Linear(__snake_case     ,			__snake_case			)

										if encoder_hid_proj_type is None:
														__magic_name__						:						Any										=None
										elif encoder_hid_proj_type == "linear":
														__magic_name__						:						List[str]										=nn.Linear(__snake_case     ,			__snake_case			)
										else:
														raise ValueError(f"unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"			)

										__magic_name__						:						Dict										=nn.Parameter(torch.zeros(1     ,			num_embeddings + additional_embeddings     ,			__snake_case			)			)

										if added_emb_type == "prd":
														__magic_name__						:						Union[str, Any]										=nn.Parameter(torch.zeros(1     ,			1     ,			__snake_case			)			)
										elif added_emb_type is None:
														__magic_name__						:						Optional[Any]										=None
										else:
														raise ValueError(
														    f"`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `'prd'` or `None`."			)

										__magic_name__						:						Optional[int]										=nn.ModuleList(
										    [
										        BasicTransformerBlock(
										            __snake_case     ,			__snake_case     ,			__snake_case     ,			dropout=__snake_case     ,			activation_fn="""gelu"""     ,			attention_bias=__snake_case     ,			)
										        for d in range(__snake_case			)
										    ]			)

										if norm_in_type == "layer":
														__magic_name__						:						Optional[Any]										=nn.LayerNorm(__snake_case			)
										elif norm_in_type is None:
														__magic_name__						:						Optional[Any]										=None
										else:
														raise ValueError(f"Unsupported norm_in_type: {norm_in_type}."			)

										__magic_name__						:						Tuple										=nn.LayerNorm(__snake_case			)

										__magic_name__						:						Union[str, Any]										=nn.Linear(__snake_case     ,			__snake_case			)

										__magic_name__						:						Any										=torch.full(
										    [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]     ,			-10000.0			)
										causal_attention_mask.triu_(1			)
										__magic_name__						:						Union[str, Any]										=causal_attention_mask[None, ...]
										self.register_buffer("""causal_attention_mask"""     ,			__snake_case     ,			persistent=__snake_case			)

										__magic_name__						:						Union[str, Any]										=nn.Parameter(torch.zeros(1     ,			__snake_case			)			)
										__magic_name__						:						Tuple										=nn.Parameter(torch.zeros(1     ,			__snake_case			)			)

						@property
						# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
						def      A__    (		self :List[str]			):
										'''simple docstring'''

										__magic_name__						:						Dict										={}

										def fn_recursive_add_processors(__snake_case :str     ,			__snake_case :torch.nn.Module     ,			__snake_case :Dict[str, AttentionProcessor]			):
														if hasattr(__snake_case     ,			"""set_processor"""			):
																		__magic_name__						:						Dict										=module.processor

														for sub_name, child in module.named_children():
																		fn_recursive_add_processors(f"{name}.{sub_name}"     ,			__snake_case     ,			__snake_case			)

														return processors

										for name, module in self.named_children():
														fn_recursive_add_processors(__snake_case     ,			__snake_case     ,			__snake_case			)

										return processors

						def      A__    (		self :Any     ,			__snake_case :Union[AttentionProcessor, Dict[str, AttentionProcessor]]			):
										'''simple docstring'''

										__magic_name__						:						Dict										=len(self.attn_processors.keys()			)

										if isinstance(__snake_case     ,			__snake_case			) and len(__snake_case			) != count:
														raise ValueError(
														    f"A dict of processors was passed, but the number of processors {len(__snake_case			)} does not match the"
														    f" number of attention layers: {count}. Please make sure to pass {count} processor classes."			)

										def fn_recursive_attn_processor(__snake_case :str     ,			__snake_case :torch.nn.Module     ,			__snake_case :List[Any]			):
														if hasattr(__snake_case     ,			"""set_processor"""			):
																		if not isinstance(__snake_case     ,			__snake_case			):
																						module.set_processor(__snake_case			)
																		else:
																						module.set_processor(processor.pop(f"{name}.processor"			)			)

														for sub_name, child in module.named_children():
																		fn_recursive_attn_processor(f"{name}.{sub_name}"     ,			__snake_case     ,			__snake_case			)

										for name, module in self.named_children():
														fn_recursive_attn_processor(__snake_case     ,			__snake_case     ,			__snake_case			)

						def      A__    (		self :str			):
										'''simple docstring'''

										self.set_attn_processor(AttnProcessor()			)

						def      A__    (		self :int     ,			__snake_case :List[Any]     ,			__snake_case :Union[torch.Tensor, float, int]     ,			__snake_case :torch.FloatTensor     ,			__snake_case :Optional[torch.FloatTensor] = None     ,			__snake_case :Optional[torch.BoolTensor] = None     ,			__snake_case :bool = True     ,			):
										'''simple docstring'''

										__magic_name__						:						Optional[int]										=hidden_states.shape[0]

										__magic_name__						:						str										=timestep
										if not torch.is_tensor(__snake_case			):
														__magic_name__						:						Dict										=torch.tensor([timesteps]     ,			dtype=torch.long     ,			device=hidden_states.device			)
										elif torch.is_tensor(__snake_case			) and len(timesteps.shape			) == 0:
														__magic_name__						:						Optional[int]										=timesteps[None].to(hidden_states.device			)

										# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
										__magic_name__						:						Optional[Any]										=timesteps * torch.ones(__snake_case     ,			dtype=timesteps.dtype     ,			device=timesteps.device			)

										__magic_name__						:						Any										=self.time_proj(__snake_case			)

										# timesteps does not contain any weights and will always return f32 tensors
										# but time_embedding might be fp16, so we need to cast here.
										__magic_name__						:						Optional[Any]										=timesteps_projected.to(dtype=self.dtype			)
										__magic_name__						:						Union[str, Any]										=self.time_embedding(__snake_case			)

										if self.embedding_proj_norm is not None:
														__magic_name__						:						int										=self.embedding_proj_norm(__snake_case			)

										__magic_name__						:						Optional[int]										=self.embedding_proj(__snake_case			)
										if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
														__magic_name__						:						Union[str, Any]										=self.encoder_hidden_states_proj(__snake_case			)
										elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
														raise ValueError("""`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set"""			)

										__magic_name__						:						Any										=self.proj_in(__snake_case			)

										__magic_name__						:						Tuple										=self.positional_embedding.to(hidden_states.dtype			)

										__magic_name__						:						Dict										=[]
										__magic_name__						:						List[str]										=0

										if encoder_hidden_states is not None:
														additional_embeds.append(__snake_case			)
														additional_embeddings_len += encoder_hidden_states.shape[1]

										if len(proj_embeddings.shape			) == 2:
														__magic_name__						:						Optional[int]										=proj_embeddings[:, None, :]

										if len(hidden_states.shape			) == 2:
														__magic_name__						:						Union[str, Any]										=hidden_states[:, None, :]

										__magic_name__						:						Optional[Any]										=additional_embeds + [
										    proj_embeddings,
										    time_embeddings[:, None, :],
										    hidden_states,
										]

										if self.prd_embedding is not None:
														__magic_name__						:						List[Any]										=self.prd_embedding.to(hidden_states.dtype			).expand(__snake_case     ,			-1     ,			-1			)
														additional_embeds.append(__snake_case			)

										__magic_name__						:						Tuple										=torch.cat(
										    __snake_case     ,			dim=1     ,			)

										# Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
										__magic_name__						:						List[Any]										=additional_embeddings_len + proj_embeddings.shape[1] + 1
										if positional_embeddings.shape[1] < hidden_states.shape[1]:
														__magic_name__						:						List[str]										=F.pad(
														    __snake_case     ,			(
														        0,
														        0,
														        additional_embeddings_len,
														        self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
														    )     ,			value=0.0     ,			)

										__magic_name__						:						str										=hidden_states + positional_embeddings

										if attention_mask is not None:
														__magic_name__						:						Tuple										=(1 - attention_mask.to(hidden_states.dtype			)) * -10000.0
														__magic_name__						:						Optional[int]										=F.pad(__snake_case     ,			(0, self.additional_embeddings)     ,			value=0.0			)
														__magic_name__						:						Any										=(attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype			)
														__magic_name__						:						Tuple										=attention_mask.repeat_interleave(self.config.num_attention_heads     ,			dim=0			)

										if self.norm_in is not None:
														__magic_name__						:						List[str]										=self.norm_in(__snake_case			)

										for block in self.transformer_blocks:
														__magic_name__						:						str										=block(__snake_case     ,			attention_mask=__snake_case			)

										__magic_name__						:						Tuple										=self.norm_out(__snake_case			)

										if self.prd_embedding is not None:
														__magic_name__						:						Optional[Any]										=hidden_states[:, -1]
										else:
														__magic_name__						:						Union[str, Any]										=hidden_states[:, additional_embeddings_len:]

										__magic_name__						:						str										=self.proj_to_clip_embeddings(__snake_case			)

										if not return_dict:
														return (predicted_image_embedding,)

										return PriorTransformerOutput(predicted_image_embedding=__snake_case			)


						def      A__    (		self :int     ,			__snake_case :Optional[int]			):
										'''simple docstring'''

										__magic_name__						:						Tuple										=(prior_latents * self.clip_std) + self.clip_mean
										return prior_latents










'''simple docstring'''





from sklearn.metrics import fa_score, matthews_corrcoef

import datasets

from .record_evaluation import evaluate as evaluate_record


UpperCamelCase_		:     Tuple          =		'''\
@article{wang2019superglue,
  title={SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems},
  author={Wang, Alex and Pruksachatkun, Yada and Nangia, Nikita and Singh, Amanpreet and Michael, Julian and Hill, Felix and Levy, Omer and Bowman, Samuel R},
  journal={arXiv preprint arXiv:1905.00537},
  year={2019}
}
'''

UpperCamelCase_		:     Optional[Any]          =		'''\
SuperGLUE (https://super.gluebenchmark.com/) is a new benchmark styled after
GLUE with a new set of more difficult language understanding tasks, improved
resources, and a new public leaderboard.
'''

UpperCamelCase_		:     Dict          =		'''
Compute SuperGLUE evaluation metric associated to each SuperGLUE dataset.
Args:
    predictions: list of predictions to score. Depending on the SuperGlUE subset:
        - for \'record\': list of question-answer dictionaries with the following keys:
            - \'idx\': index of the question as specified by the dataset
            - \'prediction_text\': the predicted answer text
        - for \'multirc\': list of question-answer dictionaries with the following keys:
            - \'idx\': index of the question-answer pair as specified by the dataset
            - \'prediction\': the predicted answer label
        - otherwise: list of predicted labels
    references: list of reference labels. Depending on the SuperGLUE subset:
        - for \'record\': list of question-answers dictionaries with the following keys:
            - \'idx\': index of the question as specified by the dataset
            - \'answers\': list of possible answers
        - otherwise: list of reference labels
Returns: depending on the SuperGLUE subset:
    - for \'record\':
        - \'exact_match\': Exact match between answer and gold answer
        - \'f1\': F1 score
    - for \'multirc\':
        - \'exact_match\': Exact match between answer and gold answer
        - \'f1_m\': Per-question macro-F1 score
        - \'f1_a\': Average F1 score over all answers
    - for \'axb\':
        \'matthews_correlation\': Matthew Correlation
    - for \'cb\':
        - \'accuracy\': Accuracy
        - \'f1\': F1 score
    - for all others:
        - \'accuracy\': Accuracy
Examples:

    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'copa\')  # any of ["copa", "rte", "wic", "wsc", "wsc.fixed", "boolq", "axg"]
    >>> predictions = [0, 1]
    >>> references = [0, 1]
    >>> results = super_glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'accuracy\': 1.0}

    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'cb\')
    >>> predictions = [0, 1]
    >>> references = [0, 1]
    >>> results = super_glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'accuracy\': 1.0, \'f1\': 1.0}

    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'record\')
    >>> predictions = [{\'idx\': {\'passage\': 0, \'query\': 0}, \'prediction_text\': \'answer\'}]
    >>> references = [{\'idx\': {\'passage\': 0, \'query\': 0}, \'answers\': [\'answer\', \'another_answer\']}]
    >>> results = super_glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'exact_match\': 1.0, \'f1\': 1.0}

    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'multirc\')
    >>> predictions = [{\'idx\': {\'answer\': 0, \'paragraph\': 0, \'question\': 0}, \'prediction\': 0}, {\'idx\': {\'answer\': 1, \'paragraph\': 2, \'question\': 3}, \'prediction\': 1}]
    >>> references = [0, 1]
    >>> results = super_glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'exact_match\': 1.0, \'f1_m\': 1.0, \'f1_a\': 1.0}

    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'axb\')
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = super_glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'matthews_correlation\': 1.0}
'''





def       __a      (  _UpperCamelCase:       Tuple   ,		_UpperCamelCase:       int       )		->					Optional[Any]:

     """simple docstring"""




     return float((preds == labels).mean()       )





def       __a      (  _UpperCamelCase:       Union[str, Any]   ,		_UpperCamelCase:       Optional[int]   ,		_UpperCamelCase:       int="binary"       )		->					Any:

     """simple docstring"""




     _snake_case				= simple_accuracy(_UpperCamelCase   ,		_UpperCamelCase       )
     _snake_case				= float(fa_score(y_true=_UpperCamelCase   ,		y_pred=_UpperCamelCase   ,		average=_UpperCamelCase       )       )
     return {
         "accuracy": acc,
         "f1": fa,
     }





def       __a      (  _UpperCamelCase:       str   ,		_UpperCamelCase:       Any       )		->					Tuple:

     """simple docstring"""




     _snake_case				= {}
     for id_pred, label in zip(_UpperCamelCase   ,		_UpperCamelCase       ):
          _snake_case				= F"""{id_pred['idx']['paragraph']}-{id_pred['idx']['question']}"""
          _snake_case				= id_pred["prediction"]
          if question_id in question_map:
               question_map[question_id].append((pred, label)       )
          else:
               _snake_case				= [(pred, label)]
     _snake_case				,   _snake_case				= [], []
     for question, preds_labels in question_map.items():
          _snake_case				,   _snake_case				= zip(*_UpperCamelCase       )
          _snake_case				= fa_score(y_true=_UpperCamelCase   ,		y_pred=_UpperCamelCase   ,		average="macro"       )
          fas.append(_UpperCamelCase       )
          _snake_case				= int(sum(pred == label for pred, label in preds_labels       ) == len(_UpperCamelCase       )       )
          ems.append(_UpperCamelCase       )
     _snake_case				= float(sum(_UpperCamelCase       ) / len(_UpperCamelCase       )       )
     _snake_case				= sum(_UpperCamelCase       ) / len(_UpperCamelCase       )
     _snake_case				= float(fa_score(y_true=_UpperCamelCase   ,		y_pred=[id_pred["prediction"] for id_pred in ids_preds]       )       )
     return {"exact_match": em, "f1_m": fa_m, "f1_a": fa_a}





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






     def 							_lowercase      (					self       )				->     Union[str, Any]:
          if self.config_name not in [
              "boolq",
              "cb",
              "copa",
              "multirc",
              "record",
              "rte",
              "wic",
              "wsc",
              "wsc.fixed",
              "axb",
              "axg",
          ]:
               raise KeyError(
                   "You should supply a configuration name selected in "
                   "[\"boolq\", \"cb\", \"copa\", \"multirc\", \"record\", \"rte\", \"wic\", \"wsc\", \"wsc.fixed\", \"axb\", \"axg\",]"       )
          return datasets.MetricInfo(
              description=_DESCRIPTION   ,citation=_CITATION   ,inputs_description=_KWARGS_DESCRIPTION   ,features=datasets.Features(self._get_feature_types()       )   ,codebase_urls=[]   ,reference_urls=[]   ,format="numpy" if not self.config_name == "record" and not self.config_name == "multirc" else None   ,)






     def 							_lowercase      (					self       )				->     Union[str, Any]:
          if self.config_name == "record":
               return {
                   "predictions": {
                       "idx": {
                           "passage": datasets.Value("int64"       ),
                           "query": datasets.Value("int64"       ),
                       },
                       "prediction_text": datasets.Value("string"       ),
                   },
                   "references": {
                       "idx": {
                           "passage": datasets.Value("int64"       ),
                           "query": datasets.Value("int64"       ),
                       },
                       "answers": datasets.Sequence(datasets.Value("string"       )       ),
                   },
               }
          elif self.config_name == "multirc":
               return {
                   "predictions": {
                       "idx": {
                           "answer": datasets.Value("int64"       ),
                           "paragraph": datasets.Value("int64"       ),
                           "question": datasets.Value("int64"       ),
                       },
                       "prediction": datasets.Value("int64"       ),
                   },
                   "references": datasets.Value("int64"       ),
               }
          else:
               return {
                   "predictions": datasets.Value("int64"       ),
                   "references": datasets.Value("int64"       ),
               }






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     Optional[Any]:
          if self.config_name == "axb":
               return {"matthews_correlation": matthews_corrcoef(_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )}
          elif self.config_name == "cb":
               return acc_and_fa(_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,fa_avg="macro"       )
          elif self.config_name == "record":
               _snake_case				= [
                   {
                       "qas": [
                           {"id": ref["idx"]["query"], "answers": [{"text": ans} for ans in ref["answers"]]}
                           for ref in references
                       ]
                   }
               ]
               _snake_case				= {pred["idx"]["query"]: pred["prediction_text"] for pred in predictions}
               return evaluate_record(_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )[0]
          elif self.config_name == "multirc":
               return evaluate_multirc(_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )
          elif self.config_name in ["copa", "rte", "wic", "wsc", "wsc.fixed", "boolq", "axg"]:
               return {"accuracy": simple_accuracy(_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )}
          else:
               raise KeyError(
                   "You should supply a configuration name selected in "
                   "[\"boolq\", \"cb\", \"copa\", \"multirc\", \"record\", \"rte\", \"wic\", \"wsc\", \"wsc.fixed\", \"axb\", \"axg\",]"       )


