


from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block




@dataclass
class  UpperCAmelCase       (							__snake_case     ):
			lowercase						       =			42




class  UpperCAmelCase       (							nn.Module     ):




			def __init__(    self				:		Any							,							__magic_name__				:		List[Any]=3							,							__magic_name__				:		int=3							,							__magic_name__				:		int=("DownEncoderBlock2D",)							,							__magic_name__				:		int=(6_4,)							,							__magic_name__				:		Optional[Any]=2							,							__magic_name__				:		List[Any]=3_2							,							__magic_name__				:		Tuple="silu"							,							__magic_name__				:		Union[str, Any]=True							,							):







						"""simple docstring"""







						super().__init__()
						UpperCamelCase								=       layers_per_block

						UpperCamelCase								=       torch.nn.Convad(
						    SCREAMING_SNAKE_CASE__							,							block_out_channels[0]							,							kernel_size=3							,							stride=1							,							padding=1							,							)

						UpperCamelCase								=       None
						UpperCamelCase								=       nn.ModuleList([]    )

						# down
						UpperCamelCase								=       block_out_channels[0]
						for i, down_block_type in enumerate(SCREAMING_SNAKE_CASE__    ):
									UpperCamelCase								=       output_channel
									UpperCamelCase								=       block_out_channels[i]
									UpperCamelCase								=       i == len(SCREAMING_SNAKE_CASE__    ) - 1

									UpperCamelCase								=       get_down_block(
									    SCREAMING_SNAKE_CASE__							,							num_layers=self.layers_per_block							,							in_channels=SCREAMING_SNAKE_CASE__							,							out_channels=SCREAMING_SNAKE_CASE__							,							add_downsample=not is_final_block							,							resnet_eps=1e-6							,							downsample_padding=0							,							resnet_act_fn=SCREAMING_SNAKE_CASE__							,							resnet_groups=SCREAMING_SNAKE_CASE__							,							attention_head_dim=SCREAMING_SNAKE_CASE__							,							temb_channels=SCREAMING_SNAKE_CASE__							,							)
									self.down_blocks.append(SCREAMING_SNAKE_CASE__    )

						# mid
						UpperCamelCase								=       UNetMidBlockaD(
						    in_channels=block_out_channels[-1]							,							resnet_eps=1e-6							,							resnet_act_fn=SCREAMING_SNAKE_CASE__							,							output_scale_factor=1							,							resnet_time_scale_shift="""default"""							,							attention_head_dim=block_out_channels[-1]							,							resnet_groups=SCREAMING_SNAKE_CASE__							,							temb_channels=SCREAMING_SNAKE_CASE__							,							)

						# out
						UpperCamelCase								=       nn.GroupNorm(num_channels=block_out_channels[-1]							,							num_groups=SCREAMING_SNAKE_CASE__							,							eps=1e-6    )
						UpperCamelCase								=       nn.SiLU()

						UpperCamelCase								=       2 * out_channels if double_z else out_channels
						UpperCamelCase								=       nn.Convad(block_out_channels[-1]							,							SCREAMING_SNAKE_CASE__							,							3							,							padding=1    )

						UpperCamelCase								=       False





			def 	lowerCamelCase_       (    self				:		Any							,							__magic_name__				:		Tuple    ):







						"""simple docstring"""







						UpperCamelCase								=       x
						UpperCamelCase								=       self.conv_in(SCREAMING_SNAKE_CASE__    )

						if self.training and self.gradient_checkpointing:

									def create_custom_forward(__magic_name__				:		List[str]    ):
												def custom_forward(*__magic_name__				:		int    ):
															return module(*SCREAMING_SNAKE_CASE__    )

												return custom_forward

									# down
									if is_torch_version(""">="""							,							"""1.11.0"""    ):
												for down_block in self.down_blocks:
															UpperCamelCase								=       torch.utils.checkpoint.checkpoint(
															    create_custom_forward(SCREAMING_SNAKE_CASE__    )							,							SCREAMING_SNAKE_CASE__							,							use_reentrant=SCREAMING_SNAKE_CASE__    )
												# middle
												UpperCamelCase								=       torch.utils.checkpoint.checkpoint(
												    create_custom_forward(self.mid_block    )							,							SCREAMING_SNAKE_CASE__							,							use_reentrant=SCREAMING_SNAKE_CASE__    )
									else:
												for down_block in self.down_blocks:
															UpperCamelCase								=       torch.utils.checkpoint.checkpoint(create_custom_forward(SCREAMING_SNAKE_CASE__    )							,							SCREAMING_SNAKE_CASE__    )
												# middle
												UpperCamelCase								=       torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block    )							,							SCREAMING_SNAKE_CASE__    )

						else:
									# down
									for down_block in self.down_blocks:
												UpperCamelCase								=       down_block(SCREAMING_SNAKE_CASE__    )

									# middle
									UpperCamelCase								=       self.mid_block(SCREAMING_SNAKE_CASE__    )

						# post-process
						UpperCamelCase								=       self.conv_norm_out(SCREAMING_SNAKE_CASE__    )
						UpperCamelCase								=       self.conv_act(SCREAMING_SNAKE_CASE__    )
						UpperCamelCase								=       self.conv_out(SCREAMING_SNAKE_CASE__    )

						return sample




class  UpperCAmelCase       (							nn.Module     ):




			def __init__(    self				:		Optional[Any]							,							__magic_name__				:		Union[str, Any]=3							,							__magic_name__				:		str=3							,							__magic_name__				:		Tuple=("UpDecoderBlock2D",)							,							__magic_name__				:		Optional[int]=(6_4,)							,							__magic_name__				:		str=2							,							__magic_name__				:		Tuple=3_2							,							__magic_name__				:		List[Any]="silu"							,							__magic_name__				:		Union[str, Any]="group"							,							):







						"""simple docstring"""







						super().__init__()
						UpperCamelCase								=       layers_per_block

						UpperCamelCase								=       nn.Convad(
						    SCREAMING_SNAKE_CASE__							,							block_out_channels[-1]							,							kernel_size=3							,							stride=1							,							padding=1							,							)

						UpperCamelCase								=       None
						UpperCamelCase								=       nn.ModuleList([]    )

						UpperCamelCase								=       in_channels if norm_type == """spatial""" else None

						# mid
						UpperCamelCase								=       UNetMidBlockaD(
						    in_channels=block_out_channels[-1]							,							resnet_eps=1e-6							,							resnet_act_fn=SCREAMING_SNAKE_CASE__							,							output_scale_factor=1							,							resnet_time_scale_shift="""default""" if norm_type == """group""" else norm_type							,							attention_head_dim=block_out_channels[-1]							,							resnet_groups=SCREAMING_SNAKE_CASE__							,							temb_channels=SCREAMING_SNAKE_CASE__							,							)

						# up
						UpperCamelCase								=       list(reversed(SCREAMING_SNAKE_CASE__    )    )
						UpperCamelCase								=       reversed_block_out_channels[0]
						for i, up_block_type in enumerate(SCREAMING_SNAKE_CASE__    ):
									UpperCamelCase								=       output_channel
									UpperCamelCase								=       reversed_block_out_channels[i]

									UpperCamelCase								=       i == len(SCREAMING_SNAKE_CASE__    ) - 1

									UpperCamelCase								=       get_up_block(
									    SCREAMING_SNAKE_CASE__							,							num_layers=self.layers_per_block + 1							,							in_channels=SCREAMING_SNAKE_CASE__							,							out_channels=SCREAMING_SNAKE_CASE__							,							prev_output_channel=SCREAMING_SNAKE_CASE__							,							add_upsample=not is_final_block							,							resnet_eps=1e-6							,							resnet_act_fn=SCREAMING_SNAKE_CASE__							,							resnet_groups=SCREAMING_SNAKE_CASE__							,							attention_head_dim=SCREAMING_SNAKE_CASE__							,							temb_channels=SCREAMING_SNAKE_CASE__							,							resnet_time_scale_shift=SCREAMING_SNAKE_CASE__							,							)
									self.up_blocks.append(SCREAMING_SNAKE_CASE__    )
									UpperCamelCase								=       output_channel

						# out
						if norm_type == "spatial":
									UpperCamelCase								=       SpatialNorm(block_out_channels[0]							,							SCREAMING_SNAKE_CASE__    )
						else:
									UpperCamelCase								=       nn.GroupNorm(num_channels=block_out_channels[0]							,							num_groups=SCREAMING_SNAKE_CASE__							,							eps=1e-6    )
						UpperCamelCase								=       nn.SiLU()
						UpperCamelCase								=       nn.Convad(block_out_channels[0]							,							SCREAMING_SNAKE_CASE__							,							3							,							padding=1    )

						UpperCamelCase								=       False





			def 	lowerCamelCase_       (    self				:		Tuple							,							__magic_name__				:		Union[str, Any]							,							__magic_name__				:		Union[str, Any]=None    ):







						"""simple docstring"""







						UpperCamelCase								=       z
						UpperCamelCase								=       self.conv_in(SCREAMING_SNAKE_CASE__    )

						UpperCamelCase								=       next(iter(self.up_blocks.parameters()    )    ).dtype
						if self.training and self.gradient_checkpointing:

									def create_custom_forward(__magic_name__				:		Optional[int]    ):
												def custom_forward(*__magic_name__				:		List[str]    ):
															return module(*SCREAMING_SNAKE_CASE__    )

												return custom_forward

									if is_torch_version(""">="""							,							"""1.11.0"""    ):
												# middle
												UpperCamelCase								=       torch.utils.checkpoint.checkpoint(
												    create_custom_forward(self.mid_block    )							,							SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__							,							use_reentrant=SCREAMING_SNAKE_CASE__    )
												UpperCamelCase								=       sample.to(SCREAMING_SNAKE_CASE__    )

												# up
												for up_block in self.up_blocks:
															UpperCamelCase								=       torch.utils.checkpoint.checkpoint(
															    create_custom_forward(SCREAMING_SNAKE_CASE__    )							,							SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__							,							use_reentrant=SCREAMING_SNAKE_CASE__    )
									else:
												# middle
												UpperCamelCase								=       torch.utils.checkpoint.checkpoint(
												    create_custom_forward(self.mid_block    )							,							SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__    )
												UpperCamelCase								=       sample.to(SCREAMING_SNAKE_CASE__    )

												# up
												for up_block in self.up_blocks:
															UpperCamelCase								=       torch.utils.checkpoint.checkpoint(create_custom_forward(SCREAMING_SNAKE_CASE__    )							,							SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__    )
						else:
									# middle
									UpperCamelCase								=       self.mid_block(SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__    )
									UpperCamelCase								=       sample.to(SCREAMING_SNAKE_CASE__    )

									# up
									for up_block in self.up_blocks:
												UpperCamelCase								=       up_block(SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__    )

        # post-process
						if latent_embeds is None:
									UpperCamelCase								=       self.conv_norm_out(SCREAMING_SNAKE_CASE__    )
						else:
									UpperCamelCase								=       self.conv_norm_out(SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__    )
						UpperCamelCase								=       self.conv_act(SCREAMING_SNAKE_CASE__    )
						UpperCamelCase								=       self.conv_out(SCREAMING_SNAKE_CASE__    )

						return sample




class  UpperCAmelCase       (							nn.Module     ):




			def __init__(    self				:		Dict							,							__magic_name__				:		Optional[Any]							,							__magic_name__				:		int							,							__magic_name__				:		Any							,							__magic_name__				:		Optional[Any]=None							,							__magic_name__				:		str="random"							,							__magic_name__				:		Union[str, Any]=False							,							__magic_name__				:		Union[str, Any]=True    ):







						"""simple docstring"""







						super().__init__()
						UpperCamelCase								=       n_e
						UpperCamelCase								=       vq_embed_dim
						UpperCamelCase								=       beta
						UpperCamelCase								=       legacy

						UpperCamelCase								=       nn.Embedding(self.n_e							,							self.vq_embed_dim    )
						self.embedding.weight.data.uniform_(-1.0 / self.n_e							,							1.0 / self.n_e    )

						UpperCamelCase								=       remap
						if self.remap is not None:
									self.register_buffer("""used"""							,							torch.tensor(np.load(self.remap    )    )    )
									UpperCamelCase								=       self.used.shape[0]
									UpperCamelCase								=       unknown_index  # "random" or "extra" or integer
									if self.unknown_index == "extra":
												UpperCamelCase								=       self.re_embed
												UpperCamelCase								=       self.re_embed + 1
									print(
									    F'Remapping {self.n_e} indices to {self.re_embed} indices. '
									    F'Using {self.unknown_index} for unknown indices.'    )
						else:
									UpperCamelCase								=       n_e

						UpperCamelCase								=       sane_index_shape




			def 	lowerCamelCase_       (    self				:		List[Any]							,							__magic_name__				:		Optional[int]    ):







						"""simple docstring"""







						UpperCamelCase								=       inds.shape
						assert len(SCREAMING_SNAKE_CASE__    ) > 1
						UpperCamelCase								=       inds.reshape(ishape[0]							,							-1    )
						UpperCamelCase								=       self.used.to(SCREAMING_SNAKE_CASE__    )
						UpperCamelCase								=       (inds[:, :, None] == used[None, None, ...]).long()
						UpperCamelCase								=       match.argmax(-1    )
						UpperCamelCase								=       match.sum(2    ) < 1
						if self.unknown_index == "random":
									UpperCamelCase								=       torch.randint(0							,							self.re_embed							,							size=new[unknown].shape    ).to(device=new.device    )
						else:
									UpperCamelCase								=       self.unknown_index
						return new.reshape(SCREAMING_SNAKE_CASE__    )




			def 	lowerCamelCase_       (    self				:		List[str]							,							__magic_name__				:		str    ):







						"""simple docstring"""







						UpperCamelCase								=       inds.shape
						assert len(SCREAMING_SNAKE_CASE__    ) > 1
						UpperCamelCase								=       inds.reshape(ishape[0]							,							-1    )
						UpperCamelCase								=       self.used.to(SCREAMING_SNAKE_CASE__    )
						if self.re_embed > self.used.shape[0]:  # extra token
									UpperCamelCase								=       0  # simply set to zero
						UpperCamelCase								=       torch.gather(used[None, :][inds.shape[0] * [0], :]							,							1							,							SCREAMING_SNAKE_CASE__    )
						return back.reshape(SCREAMING_SNAKE_CASE__    )




			def 	lowerCamelCase_       (    self				:		Any							,							__magic_name__				:		Optional[int]    ):







						"""simple docstring"""







						UpperCamelCase								=       z.permute(0							,							2							,							3							,							1    ).contiguous()
						UpperCamelCase								=       z.view(-1							,							self.vq_embed_dim    )

						# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
						UpperCamelCase								=       torch.argmin(torch.cdist(SCREAMING_SNAKE_CASE__							,							self.embedding.weight    )							,							dim=1    )

						UpperCamelCase								=       self.embedding(SCREAMING_SNAKE_CASE__    ).view(z.shape    )
						UpperCamelCase								=       None
						UpperCamelCase								=       None

						# compute loss for embedding
						if not self.legacy:
									UpperCamelCase								=       self.beta * torch.mean((z_q.detach() - z) ** 2    ) + torch.mean((z_q - z.detach()) ** 2    )
						else:
									UpperCamelCase								=       torch.mean((z_q.detach() - z) ** 2    ) + self.beta * torch.mean((z_q - z.detach()) ** 2    )

						# preserve gradients
						UpperCamelCase								=       z + (z_q - z).detach()

						# reshape back to match original input shape
						UpperCamelCase								=       z_q.permute(0							,							3							,							1							,							2    ).contiguous()

						if self.remap is not None:
									UpperCamelCase								=       min_encoding_indices.reshape(z.shape[0]							,							-1    )  # add batch axis
									UpperCamelCase								=       self.remap_to_used(SCREAMING_SNAKE_CASE__    )
									UpperCamelCase								=       min_encoding_indices.reshape(-1							,							1    )  # flatten

						if self.sane_index_shape:
									UpperCamelCase								=       min_encoding_indices.reshape(z_q.shape[0]							,							z_q.shape[2]							,							z_q.shape[3]    )

						return z_q, loss, (perplexity, min_encodings, min_encoding_indices)





			def 	lowerCamelCase_       (    self				:		Dict							,							__magic_name__				:		Any							,							__magic_name__				:		List[Any]    ):







						"""simple docstring"""







						if self.remap is not None:
									UpperCamelCase								=       indices.reshape(shape[0]							,							-1    )  # add batch axis
									UpperCamelCase								=       self.unmap_to_all(SCREAMING_SNAKE_CASE__    )
									UpperCamelCase								=       indices.reshape(-1    )  # flatten again

						# get quantized latent vectors
						UpperCamelCase								=       self.embedding(SCREAMING_SNAKE_CASE__    )

						if shape is not None:
									UpperCamelCase								=       z_q.view(SCREAMING_SNAKE_CASE__    )
									# reshape back to match original input shape
									UpperCamelCase								=       z_q.permute(0							,							3							,							1							,							2    ).contiguous()

						return z_q




class  UpperCAmelCase       (							__snake_case     ):




			def __init__(    self				:		Optional[Any]							,							__magic_name__				:		Any							,							__magic_name__				:		Optional[int]=False    ):







						"""simple docstring"""







						UpperCamelCase								=       parameters
						UpperCamelCase	,	UpperCamelCase								=       torch.chunk(SCREAMING_SNAKE_CASE__							,							2							,							dim=1    )
						UpperCamelCase								=       torch.clamp(self.logvar							,							-30.0							,							20.0    )
						UpperCamelCase								=       deterministic
						UpperCamelCase								=       torch.exp(0.5 * self.logvar    )
						UpperCamelCase								=       torch.exp(self.logvar    )
						if self.deterministic:
									UpperCamelCase								=       UpperCamelCase								=       torch.zeros_like(
									    self.mean							,							device=self.parameters.device							,							dtype=self.parameters.dtype    )




			def 	lowerCamelCase_       (    self				:		Any							,							__magic_name__				:		Optional[torch.Generator] = None    ):







						"""simple docstring"""







						UpperCamelCase								=       randn_tensor(
						    self.mean.shape							,							generator=SCREAMING_SNAKE_CASE__							,							device=self.parameters.device							,							dtype=self.parameters.dtype    )
						UpperCamelCase								=       self.mean + self.std * sample
						return x




			def 	lowerCamelCase_       (    self				:		str							,							__magic_name__				:		List[str]=None    ):







						"""simple docstring"""







						if self.deterministic:
									return torch.Tensor([0.0]    )
						else:
									if other is None:
												return 0.5 * torch.sum(torch.pow(self.mean							,							2    ) + self.var - 1.0 - self.logvar							,							dim=[1, 2, 3]    )
									else:
												return 0.5 * torch.sum(
												    torch.pow(self.mean - other.mean							,							2    ) / other.var
												    + self.var / other.var
												    - 1.0
												    - self.logvar
												    + other.logvar							,							dim=[1, 2, 3]							,							)




			def 	lowerCamelCase_       (    self				:		int							,							__magic_name__				:		int							,							__magic_name__				:		Union[str, Any]=[1, 2, 3]    ):







						"""simple docstring"""







						if self.deterministic:
									return torch.Tensor([0.0]    )
						UpperCamelCase								=       np.log(2.0 * np.pi    )
						return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean							,							2    ) / self.var							,							dim=SCREAMING_SNAKE_CASE__    )





			def 	lowerCamelCase_       (    self				:		int    ):







						"""simple docstring"""







						return self.mean






from collections import OrderedDict
from typing import Mapping

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


_snake_case      = logging.get_logger(__name__)

_snake_case      = {
    '''roberta-base''': '''https://huggingface.co/roberta-base/resolve/main/config.json''',
    '''roberta-large''': '''https://huggingface.co/roberta-large/resolve/main/config.json''',
    '''roberta-large-mnli''': '''https://huggingface.co/roberta-large-mnli/resolve/main/config.json''',
    '''distilroberta-base''': '''https://huggingface.co/distilroberta-base/resolve/main/config.json''',
    '''roberta-base-openai-detector''': '''https://huggingface.co/roberta-base-openai-detector/resolve/main/config.json''',
    '''roberta-large-openai-detector''': '''https://huggingface.co/roberta-large-openai-detector/resolve/main/config.json''',
}
class       _lowerCAmelCase  (		__magic_name__							):

 """simple docstring"""
 SCREAMING_SNAKE_CASE_				:					Dict       						="roberta"

 def __init__(						self			:			Union[str, Any]					, SCREAMING_SNAKE_CASE__			:			Tuple=5_02_65					, SCREAMING_SNAKE_CASE__			:			List[str]=7_68					, SCREAMING_SNAKE_CASE__			:			List[str]=12					, SCREAMING_SNAKE_CASE__			:			Any=12					, SCREAMING_SNAKE_CASE__			:			Tuple=30_72					, SCREAMING_SNAKE_CASE__			:			str="gelu"					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]=0.1					, SCREAMING_SNAKE_CASE__			:			int=0.1					, SCREAMING_SNAKE_CASE__			:			int=5_12					, SCREAMING_SNAKE_CASE__			:			Optional[Any]=2					, SCREAMING_SNAKE_CASE__			:			Dict=0.02					, SCREAMING_SNAKE_CASE__			:			List[str]=1e-12					, SCREAMING_SNAKE_CASE__			:			Any=1					, SCREAMING_SNAKE_CASE__			:			Optional[Any]=0					, SCREAMING_SNAKE_CASE__			:			str=2					, SCREAMING_SNAKE_CASE__			:			Any="absolute"					, SCREAMING_SNAKE_CASE__			:			Any=True					, SCREAMING_SNAKE_CASE__			:			List[Any]=None					, **SCREAMING_SNAKE_CASE__			:			Optional[int]					, ):






      """simple docstring"""



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

      UpperCamelCase						    =					vocab_size
      UpperCamelCase						    =					hidden_size
      UpperCamelCase						    =					num_hidden_layers
      UpperCamelCase						    =					num_attention_heads
      UpperCamelCase						    =					hidden_act
      UpperCamelCase						    =					intermediate_size
      UpperCamelCase						    =					hidden_dropout_prob
      UpperCamelCase						    =					attention_probs_dropout_prob
      UpperCamelCase						    =					max_position_embeddings
      UpperCamelCase						    =					type_vocab_size
      UpperCamelCase						    =					initializer_range
      UpperCamelCase						    =					layer_norm_eps
      UpperCamelCase						    =					position_embedding_type
      UpperCamelCase						    =					use_cache
      UpperCamelCase						    =					classifier_dropout




class       _lowerCAmelCase  (		__magic_name__							):

 """simple docstring"""

 @property
 def      __lowerCAmelCase  (						self			:			Optional[int]   ):






      """simple docstring"""



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


