


'''simple docstring'''







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 	UpperCamelCase__					( nn.Module	):



     """simple docstring"""






     def __init__( self      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__=0.0      ,	snake_case__ = None      ,	snake_case__ = "geglu"      ,	snake_case__ = None      ,	snake_case__ = False      ,	snake_case__ = False      ,	snake_case__ = False      ,	snake_case__ = False      ,	snake_case__ = True      ,	snake_case__ = "layer_norm"      ,	snake_case__ = False      ,	):

            '''simple docstring'''

            super().__init__()
            _lowerCAmelCase    :							Any   			=       only_cross_attention

            _lowerCAmelCase    :							Optional[Any]   			=       (num_embeds_ada_norm is not None) and norm_type == 'ada_norm_zero'
            _lowerCAmelCase    :							int   			=       (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:
                   _lowerCAmelCase    :							Any   			=       AdaLayerNorm(snake_case__      ,	snake_case__   )
            elif self.use_ada_layer_norm_zero:
                   _lowerCAmelCase    :							str   			=       AdaLayerNormZero(snake_case__      ,	snake_case__   )
            else:
                   _lowerCAmelCase    :							Tuple   			=       nn.LayerNorm(snake_case__      ,	elementwise_affine=snake_case__   )
            _lowerCAmelCase    :							Dict   			=       Attention(
                query_dim=snake_case__      ,	heads=snake_case__      ,	dim_head=snake_case__      ,	dropout=snake_case__      ,	bias=snake_case__      ,	cross_attention_dim=cross_attention_dim if only_cross_attention else None      ,	upcast_attention=snake_case__      ,	)

            # 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.
                   _lowerCAmelCase    :							Dict   			=       (
                       AdaLayerNorm(snake_case__      ,	snake_case__   )
                       if self.use_ada_layer_norm
                       else nn.LayerNorm(snake_case__      ,	elementwise_affine=snake_case__   )
                   )
                   _lowerCAmelCase    :							Any   			=       Attention(
                       query_dim=snake_case__      ,	cross_attention_dim=cross_attention_dim if not double_self_attention else None      ,	heads=snake_case__      ,	dim_head=snake_case__      ,	dropout=snake_case__      ,	bias=snake_case__      ,	upcast_attention=snake_case__      ,	)  # is self-attn if encoder_hidden_states is none
            else:
                   _lowerCAmelCase    :							List[Any]   			=       None
                   _lowerCAmelCase    :							Dict   			=       None

            # 3. Feed-forward
            _lowerCAmelCase    :							Tuple   			=       nn.LayerNorm(snake_case__      ,	elementwise_affine=snake_case__   )
            _lowerCAmelCase    :							int   			=       FeedForward(snake_case__      ,	dropout=snake_case__      ,	activation_fn=snake_case__      ,	final_dropout=snake_case__   )

            # let chunk size default to None
            _lowerCAmelCase    :							int   			=       None
            _lowerCAmelCase    :							Optional[Any]   			=       0







     def a	( self      ,	snake_case__      ,	snake_case__   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Any   			=       chunk_size
            _lowerCAmelCase    :							List[str]   			=       dim







     def a	( self      ,	snake_case__      ,	snake_case__ = None      ,	snake_case__ = None      ,	snake_case__ = None      ,	snake_case__ = None      ,	snake_case__ = None      ,	snake_case__ = None      ,	):

            '''simple docstring'''

            if self.use_ada_layer_norm:
                   _lowerCAmelCase    :							Optional[Any]   			=       self.norma(snake_case__      ,	snake_case__   )
            elif self.use_ada_layer_norm_zero:
                   _lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase    :							Tuple   			=       self.norma(
                       snake_case__      ,	snake_case__      ,	snake_case__      ,	hidden_dtype=hidden_states.dtype   )
            else:
                   _lowerCAmelCase    :							Optional[int]   			=       self.norma(snake_case__   )

            _lowerCAmelCase    :							Any   			=       cross_attention_kwargs if cross_attention_kwargs is not None else {}

            _lowerCAmelCase    :							Any   			=       self.attna(
                snake_case__      ,	encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None      ,	attention_mask=snake_case__      ,	**snake_case__      ,	)
            if self.use_ada_layer_norm_zero:
                   _lowerCAmelCase    :							int   			=       gate_msa.unsqueeze(1   ) * attn_output
            _lowerCAmelCase    :							List[str]   			=       attn_output + hidden_states

            # 2. Cross-Attention
            if self.attna is not None:
                   _lowerCAmelCase    :							Optional[int]   			=       (
                       self.norma(snake_case__      ,	snake_case__   ) if self.use_ada_layer_norm else self.norma(snake_case__   )
                   )

                   _lowerCAmelCase    :							Optional[int]   			=       self.attna(
                       snake_case__      ,	encoder_hidden_states=snake_case__      ,	attention_mask=snake_case__      ,	**snake_case__      ,	)
                   _lowerCAmelCase    :							List[str]   			=       attn_output + hidden_states

            # 3. Feed-forward
            _lowerCAmelCase    :							str   			=       self.norma(snake_case__   )

            if self.use_ada_layer_norm_zero:
                   _lowerCAmelCase    :							Optional[Any]   			=       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`.'   )

                   _lowerCAmelCase    :							List[Any]   			=       norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
                   _lowerCAmelCase    :							str   			=       torch.cat(
                       [self.ff(snake_case__   ) for hid_slice in norm_hidden_states.chunk(snake_case__      ,	dim=self._chunk_dim   )]      ,	dim=self._chunk_dim      ,	)
            else:
                   _lowerCAmelCase    :							Union[str, Any]   			=       self.ff(snake_case__   )

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

            _lowerCAmelCase    :							str   			=       ff_output + hidden_states

            return hidden_states
class 	UpperCamelCase__					( nn.Module	):



     """simple docstring"""






     def __init__( self      ,	snake_case__      ,	snake_case__ = None      ,	snake_case__ = 4      ,	snake_case__ = 0.0      ,	snake_case__ = "geglu"      ,	snake_case__ = False      ,	):

            '''simple docstring'''

            super().__init__()
            _lowerCAmelCase    :							int   			=       int(dim * mult   )
            _lowerCAmelCase    :							Optional[Any]   			=       dim_out if dim_out is not None else dim

            if activation_fn == "gelu":
                   _lowerCAmelCase    :							List[str]   			=       GELU(snake_case__      ,	snake_case__   )
            if activation_fn == "gelu-approximate":
                   _lowerCAmelCase    :							List[str]   			=       GELU(snake_case__      ,	snake_case__      ,	approximate='tanh'   )
            elif activation_fn == "geglu":
                   _lowerCAmelCase    :							str   			=       GEGLU(snake_case__      ,	snake_case__   )
            elif activation_fn == "geglu-approximate":
                   _lowerCAmelCase    :							List[Any]   			=       ApproximateGELU(snake_case__      ,	snake_case__   )

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







     def a	( self      ,	snake_case__   ):

            '''simple docstring'''

            for module in self.net:
                   _lowerCAmelCase    :							List[Any]   			=       module(snake_case__   )
            return hidden_states
class 	UpperCamelCase__					( nn.Module	):



     """simple docstring"""






     def __init__( self      ,	snake_case__      ,	snake_case__      ,	snake_case__ = "none"   ):

            '''simple docstring'''

            super().__init__()
            _lowerCAmelCase    :							Dict   			=       nn.Linear(snake_case__      ,	snake_case__   )
            _lowerCAmelCase    :							str   			=       approximate







     def a	( self      ,	snake_case__   ):

            '''simple docstring'''

            if gate.device.type != "mps":
                   return F.gelu(snake_case__      ,	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      ,	snake_case__   ):

            '''simple docstring'''

            _lowerCAmelCase    :							int   			=       self.proj(snake_case__   )
            _lowerCAmelCase    :							List[str]   			=       self.gelu(snake_case__   )
            return hidden_states
class 	UpperCamelCase__					( nn.Module	):



     """simple docstring"""






     def __init__( self      ,	snake_case__      ,	snake_case__   ):

            '''simple docstring'''

            super().__init__()
            _lowerCAmelCase    :							Union[str, Any]   			=       nn.Linear(snake_case__      ,	dim_out * 2   )







     def a	( self      ,	snake_case__   ):

            '''simple docstring'''

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







     def a	( self      ,	snake_case__   ):

            '''simple docstring'''

            _lowerCAmelCase			,     _lowerCAmelCase    :							Union[str, Any]   			=       self.proj(snake_case__   ).chunk(2      ,	dim=-1   )
            return hidden_states * self.gelu(snake_case__   )
class 	UpperCamelCase__					( nn.Module	):



     """simple docstring"""






     def __init__( self      ,	snake_case__      ,	snake_case__   ):

            '''simple docstring'''

            super().__init__()
            _lowerCAmelCase    :							List[Any]   			=       nn.Linear(snake_case__      ,	snake_case__   )







     def a	( self      ,	snake_case__   ):

            '''simple docstring'''

            _lowerCAmelCase    :							str   			=       self.proj(snake_case__   )
            return x * torch.sigmoid(1.702 * x   )
class 	UpperCamelCase__					( nn.Module	):



     """simple docstring"""






     def __init__( self      ,	snake_case__      ,	snake_case__   ):

            '''simple docstring'''

            super().__init__()
            _lowerCAmelCase    :							List[str]   			=       nn.Embedding(snake_case__      ,	snake_case__   )
            _lowerCAmelCase    :							Optional[int]   			=       nn.SiLU()
            _lowerCAmelCase    :							Dict   			=       nn.Linear(snake_case__      ,	embedding_dim * 2   )
            _lowerCAmelCase    :							Any   			=       nn.LayerNorm(snake_case__      ,	elementwise_affine=snake_case__   )







     def a	( self      ,	snake_case__      ,	snake_case__   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Dict   			=       self.linear(self.silu(self.emb(snake_case__   )   )   )
            _lowerCAmelCase			,     _lowerCAmelCase    :							List[Any]   			=       torch.chunk(snake_case__      ,	2   )
            _lowerCAmelCase    :							List[str]   			=       self.norm(snake_case__   ) * (1 + scale) + shift
            return x
class 	UpperCamelCase__					( nn.Module	):



     """simple docstring"""






     def __init__( self      ,	snake_case__      ,	snake_case__   ):

            '''simple docstring'''

            super().__init__()

            _lowerCAmelCase    :							str   			=       CombinedTimestepLabelEmbeddings(snake_case__      ,	snake_case__   )

            _lowerCAmelCase    :							List[Any]   			=       nn.SiLU()
            _lowerCAmelCase    :							Dict   			=       nn.Linear(snake_case__      ,	6 * embedding_dim      ,	bias=snake_case__   )
            _lowerCAmelCase    :							int   			=       nn.LayerNorm(snake_case__      ,	elementwise_affine=snake_case__      ,	eps=1E-6   )







     def a	( self      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__=None   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Optional[int]   			=       self.linear(self.silu(self.emb(snake_case__      ,	snake_case__      ,	hidden_dtype=snake_case__   )   )   )
            _lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase    :							Union[str, Any]   			=       emb.chunk(6      ,	dim=1   )
            _lowerCAmelCase    :							Optional[int]   			=       self.norm(snake_case__   ) * (1 + scale_msa[:, None]) + shift_msa[:, None]
            return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
class 	UpperCamelCase__					( nn.Module	):



     """simple docstring"""






     def __init__( self      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__ = None      ,	snake_case__ = 1E-5   ):

            '''simple docstring'''

            super().__init__()
            _lowerCAmelCase    :							Optional[int]   			=       num_groups
            _lowerCAmelCase    :							Any   			=       eps

            if act_fn is None:
                   _lowerCAmelCase    :							Dict   			=       None
            else:
                   _lowerCAmelCase    :							Optional[int]   			=       get_activation(snake_case__   )

            _lowerCAmelCase    :							List[str]   			=       nn.Linear(snake_case__      ,	out_dim * 2   )







     def a	( self      ,	snake_case__      ,	snake_case__   ):

            '''simple docstring'''

            if self.act:
                   _lowerCAmelCase    :							Union[str, Any]   			=       self.act(snake_case__   )
            _lowerCAmelCase    :							str   			=       self.linear(snake_case__   )
            _lowerCAmelCase    :							str   			=       emb[:, :, None, None]
            _lowerCAmelCase			,     _lowerCAmelCase    :							List[str]   			=       emb.chunk(2      ,	dim=1   )

            _lowerCAmelCase    :							List[str]   			=       F.group_norm(snake_case__      ,	self.num_groups      ,	eps=self.eps   )
            _lowerCAmelCase    :							List[str]   			=       x * (1 + scale) + shift
            return x





'''simple docstring'''







from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


lowerCAmelCase    :     str					=						logging.get_logger(__name__)

lowerCAmelCase    :     int					=						{
    """facebook/data2vec-text-base""": """https://huggingface.co/data2vec/resolve/main/config.json""",
}
class 	UpperCamelCase__					( SCREAMING_SNAKE_CASE_	):



     """simple docstring"""






     __magic_name__    			=    "data2vec-text"







     def __init__( self      ,	snake_case__=3_0522      ,	snake_case__=768      ,	snake_case__=12      ,	snake_case__=12      ,	snake_case__=3072      ,	snake_case__="gelu"      ,	snake_case__=0.1      ,	snake_case__=0.1      ,	snake_case__=512      ,	snake_case__=2      ,	snake_case__=0.02      ,	snake_case__=1E-12      ,	snake_case__=1      ,	snake_case__=0      ,	snake_case__=2      ,	snake_case__="absolute"      ,	snake_case__=True      ,	snake_case__=None      ,	**snake_case__      ,	):

            '''simple docstring'''

            super().__init__(pad_token_id=snake_case__      ,	bos_token_id=snake_case__      ,	eos_token_id=snake_case__      ,	**snake_case__   )

            _lowerCAmelCase    :							List[Any]   			=       vocab_size
            _lowerCAmelCase    :							Tuple   			=       hidden_size
            _lowerCAmelCase    :							Dict   			=       num_hidden_layers
            _lowerCAmelCase    :							int   			=       num_attention_heads
            _lowerCAmelCase    :							str   			=       hidden_act
            _lowerCAmelCase    :							Any   			=       intermediate_size
            _lowerCAmelCase    :							Any   			=       hidden_dropout_prob
            _lowerCAmelCase    :							Optional[int]   			=       attention_probs_dropout_prob
            _lowerCAmelCase    :							str   			=       max_position_embeddings
            _lowerCAmelCase    :							Any   			=       type_vocab_size
            _lowerCAmelCase    :							int   			=       initializer_range
            _lowerCAmelCase    :							List[str]   			=       layer_norm_eps
            _lowerCAmelCase    :							List[Any]   			=       position_embedding_type
            _lowerCAmelCase    :							str   			=       use_cache
            _lowerCAmelCase    :							Union[str, Any]   			=       classifier_dropout
class 	UpperCamelCase__					( SCREAMING_SNAKE_CASE_	):



     """simple docstring"""






     @property
     def a	( self   ):

            '''simple docstring'''

            if self.task == "multiple-choice":
                   _lowerCAmelCase    :							Union[str, Any]   			=       {0: 'batch', 1: 'choice', 2: 'sequence'}
            else:
                   _lowerCAmelCase    :							List[Any]   			=       {0: 'batch', 1: 'sequence'}
            return OrderedDict(
                [
                    ('input_ids', dynamic_axis),
                    ('attention_mask', dynamic_axis),
                ]   )

