


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

		"""simple docstring"""

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

							'''simple docstring'''


							super().__init__()
							UpperCamelCase__   :Dict												= only_cross_attention

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

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

							# 3. Feed-forward
							UpperCamelCase__   :List[Any]												= nn.LayerNorm(UpperCamelCase_    ,   elementwise_affine=UpperCamelCase_							)
							UpperCamelCase__   :str												= FeedForward(UpperCamelCase_    ,   dropout=UpperCamelCase_    ,   activation_fn=UpperCamelCase_    ,   final_dropout=UpperCamelCase_							)

							# let chunk size default to None
							UpperCamelCase__   :List[Any]												= None
							UpperCamelCase__   :Any												= 0




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__   :Any												= chunk_size
							UpperCamelCase__   :str												= dim




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_ = None    ,   UpperCamelCase_ = None    ,   UpperCamelCase_ = None    ,   UpperCamelCase_ = None    ,   UpperCamelCase_ = None    ,   UpperCamelCase_ = None    ,   ):

							'''simple docstring'''


							if self.use_ada_layer_norm:
												UpperCamelCase__   :int												= self.norma(UpperCamelCase_    ,   UpperCamelCase_							)
							elif self.use_ada_layer_norm_zero:
												UpperCamelCase__					,				UpperCamelCase__					,				UpperCamelCase__					,				UpperCamelCase__					,				UpperCamelCase__   :Union[str, Any]												= self.norma(
												    UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   hidden_dtype=hidden_states.dtype							)
							else:
												UpperCamelCase__   :Any												= self.norma(UpperCamelCase_							)

							UpperCamelCase__   :List[str]												= cross_attention_kwargs if cross_attention_kwargs is not None else {}

							UpperCamelCase__   :int												= self.attna(
							    UpperCamelCase_    ,   encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None    ,   attention_mask=UpperCamelCase_    ,   **UpperCamelCase_    ,   )
							if self.use_ada_layer_norm_zero:
												UpperCamelCase__   :int												= gate_msa.unsqueeze(1							) * attn_output
							UpperCamelCase__   :Optional[int]												= attn_output + hidden_states

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

												UpperCamelCase__   :Any												= self.attna(
												    UpperCamelCase_    ,   encoder_hidden_states=UpperCamelCase_    ,   attention_mask=UpperCamelCase_    ,   **UpperCamelCase_    ,   )
												UpperCamelCase__   :Tuple												= attn_output + hidden_states

							# 3. Feed-forward
							UpperCamelCase__   :List[Any]												= self.norma(UpperCamelCase_							)

							if self.use_ada_layer_norm_zero:
												UpperCamelCase__   :str												= 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`.'''							)

												UpperCamelCase__   :List[str]												= norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
												UpperCamelCase__   :Tuple												= torch.cat(
												    [self.ff(UpperCamelCase_							) for hid_slice in norm_hidden_states.chunk(UpperCamelCase_    ,   dim=self._chunk_dim							)]    ,   dim=self._chunk_dim    ,   )
							else:
												UpperCamelCase__   :Optional[Any]												= self.ff(UpperCamelCase_							)

							if self.use_ada_layer_norm_zero:
												UpperCamelCase__   :Union[str, Any]												= gate_mlp.unsqueeze(1							) * ff_output

							UpperCamelCase__   :Optional[Any]												= ff_output + hidden_states

							return hidden_states



class    lowercase	(							nn.Module					):

		"""simple docstring"""

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

							'''simple docstring'''


							super().__init__()
							UpperCamelCase__   :Tuple												= int(dim * mult							)
							UpperCamelCase__   :List[str]												= dim_out if dim_out is not None else dim

							if activation_fn == "gelu":
												UpperCamelCase__   :Optional[int]												= GELU(UpperCamelCase_    ,   UpperCamelCase_							)
							if activation_fn == "gelu-approximate":
												UpperCamelCase__   :Any												= GELU(UpperCamelCase_    ,   UpperCamelCase_    ,   approximate='''tanh'''							)
							elif activation_fn == "geglu":
												UpperCamelCase__   :int												= GEGLU(UpperCamelCase_    ,   UpperCamelCase_							)
							elif activation_fn == "geglu-approximate":
												UpperCamelCase__   :Tuple												= ApproximateGELU(UpperCamelCase_    ,   UpperCamelCase_							)

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




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_							):

							'''simple docstring'''


							for module in self.net:
												UpperCamelCase__   :Optional[int]												= module(UpperCamelCase_							)
							return hidden_states



class    lowercase	(							nn.Module					):

		"""simple docstring"""

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

							'''simple docstring'''


							super().__init__()
							UpperCamelCase__   :Any												= nn.Linear(UpperCamelCase_    ,   UpperCamelCase_							)
							UpperCamelCase__   :int												= approximate




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_							):

							'''simple docstring'''


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

							'''simple docstring'''


							UpperCamelCase__   :List[str]												= self.proj(UpperCamelCase_							)
							UpperCamelCase__   :Optional[int]												= self.gelu(UpperCamelCase_							)
							return hidden_states



class    lowercase	(							nn.Module					):

		"""simple docstring"""

		def __init__(		self    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							super().__init__()
							UpperCamelCase__   :Optional[Any]												= nn.Linear(UpperCamelCase_    ,   dim_out * 2							)




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_							):

							'''simple docstring'''


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




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__					,				UpperCamelCase__   :Any												= self.proj(UpperCamelCase_							).chunk(2    ,   dim=-1							)
							return hidden_states * self.gelu(UpperCamelCase_							)



class    lowercase	(							nn.Module					):

		"""simple docstring"""

		def __init__(		self    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							super().__init__()
							UpperCamelCase__   :List[str]												= nn.Linear(UpperCamelCase_    ,   UpperCamelCase_							)




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__   :Optional[int]												= self.proj(UpperCamelCase_							)
							return x * torch.sigmoid(1.702 * x							)



class    lowercase	(							nn.Module					):

		"""simple docstring"""

		def __init__(		self    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							super().__init__()
							UpperCamelCase__   :Union[str, Any]												= nn.Embedding(UpperCamelCase_    ,   UpperCamelCase_							)
							UpperCamelCase__   :List[str]												= nn.SiLU()
							UpperCamelCase__   :Union[str, Any]												= nn.Linear(UpperCamelCase_    ,   embedding_dim * 2							)
							UpperCamelCase__   :Optional[Any]												= nn.LayerNorm(UpperCamelCase_    ,   elementwise_affine=UpperCamelCase_							)




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__   :Optional[Any]												= self.linear(self.silu(self.emb(UpperCamelCase_							)							)							)
							UpperCamelCase__					,				UpperCamelCase__   :List[str]												= torch.chunk(UpperCamelCase_    ,   2							)
							UpperCamelCase__   :Union[str, Any]												= self.norm(UpperCamelCase_							) * (1 + scale) + shift
							return x



class    lowercase	(							nn.Module					):

		"""simple docstring"""

		def __init__(		self    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							super().__init__()

							UpperCamelCase__   :List[str]												= CombinedTimestepLabelEmbeddings(UpperCamelCase_    ,   UpperCamelCase_							)

							UpperCamelCase__   :int												= nn.SiLU()
							UpperCamelCase__   :Tuple												= nn.Linear(UpperCamelCase_    ,   6 * embedding_dim    ,   bias=UpperCamelCase_							)
							UpperCamelCase__   :str												= nn.LayerNorm(UpperCamelCase_    ,   elementwise_affine=UpperCamelCase_    ,   eps=1e-6							)




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_=None							):

							'''simple docstring'''


							UpperCamelCase__   :int												= self.linear(self.silu(self.emb(UpperCamelCase_    ,   UpperCamelCase_    ,   hidden_dtype=UpperCamelCase_							)							)							)
							UpperCamelCase__					,				UpperCamelCase__					,				UpperCamelCase__					,				UpperCamelCase__					,				UpperCamelCase__					,				UpperCamelCase__   :Dict												= emb.chunk(6    ,   dim=1							)
							UpperCamelCase__   :List[str]												= self.norm(UpperCamelCase_							) * (1 + scale_msa[:, None]) + shift_msa[:, None]
							return x, gate_msa, shift_mlp, scale_mlp, gate_mlp



class    lowercase	(							nn.Module					):

		"""simple docstring"""

		def __init__(		self    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_ = None    ,   UpperCamelCase_ = 1e-5							):

							'''simple docstring'''


							super().__init__()
							UpperCamelCase__   :List[Any]												= num_groups
							UpperCamelCase__   :Optional[Any]												= eps

							if act_fn is None:
												UpperCamelCase__   :Tuple												= None
							else:
												UpperCamelCase__   :List[str]												= get_activation(UpperCamelCase_							)

							UpperCamelCase__   :Dict												= nn.Linear(UpperCamelCase_    ,   out_dim * 2							)




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							if self.act:
												UpperCamelCase__   :int												= self.act(UpperCamelCase_							)
							UpperCamelCase__   :Union[str, Any]												= self.linear(UpperCamelCase_							)
							UpperCamelCase__   :str												= emb[:, :, None, None]
							UpperCamelCase__					,				UpperCamelCase__   :str												= emb.chunk(2    ,   dim=1							)

							UpperCamelCase__   :str												= F.group_norm(UpperCamelCase_    ,   self.num_groups    ,   eps=self.eps							)
							UpperCamelCase__   :Optional[Any]												= x * (1 + scale) + shift
							return x



'''simple docstring'''
from typing import List, Optional

from tokenizers import ByteLevelBPETokenizer

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_blenderbot_small import BlenderbotSmallTokenizer


__snake_case							      =     logging.get_logger(__name__)

__snake_case							      =     {
    '''vocab_file''': '''vocab.json''',
    '''merges_file''': '''merges.txt''',
    '''tokenizer_config_file''': '''tokenizer_config.json''',
}

__snake_case							      =     {
    '''vocab_file''': {
        '''facebook/blenderbot_small-90M''': '''https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/vocab.json'''
    },
    '''merges_file''': {
        '''facebook/blenderbot_small-90M''': '''https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/merges.txt'''
    },
    '''tokenizer_config_file''': {
        '''facebook/blenderbot_small-90M''': (
            '''https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/tokenizer_config.json'''
        )
    },
}

__snake_case							      =     {
    '''facebook/blenderbot_small-90M''': 512,
}
class    lowercase	(							A__					):

		"""simple docstring"""

		_a     =				VOCAB_FILES_NAMES
		_a     =				PRETRAINED_VOCAB_FILES_MAP
		_a     =				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
		_a     =				BlenderbotSmallTokenizer




		def __init__(		self    ,   UpperCamelCase_=None    ,   UpperCamelCase_=None    ,   UpperCamelCase_="<|endoftext|>"    ,   UpperCamelCase_="<|endoftext|>"    ,   UpperCamelCase_="<|endoftext|>"    ,   UpperCamelCase_=False    ,   UpperCamelCase_=True    ,   **UpperCamelCase_    ,   ):

							'''simple docstring'''


							super().__init__(
							    ByteLevelBPETokenizer(
							        vocab=UpperCamelCase_    ,   merges=UpperCamelCase_    ,   add_prefix_space=UpperCamelCase_    ,   trim_offsets=UpperCamelCase_    ,   )    ,   bos_token=UpperCamelCase_    ,   eos_token=UpperCamelCase_    ,   unk_token=UpperCamelCase_    ,   **UpperCamelCase_    ,   )
							UpperCamelCase__   :Union[str, Any]												= add_prefix_space




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_=None							):

							'''simple docstring'''


							UpperCamelCase__   :List[Any]												= [self.bos_token_id] + token_ids_a + [self.eos_token_id]
							if token_ids_a is None:
												return output

							return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_ = None							):

							'''simple docstring'''


							UpperCamelCase__   :Optional[int]												= [self.sep_token_id]
							UpperCamelCase__   :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]