






'''simple docstring'''



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 							__SCREAMING_SNAKE_CASE       (							__lowercase     ):
      __a  					=42







class 							__SCREAMING_SNAKE_CASE       (							nn.Module     ):




      def __init__(  self ,   lowerCamelCase=3 ,   lowerCamelCase=3 ,   lowerCamelCase=("DownEncoderBlock2D",) ,   lowerCamelCase=(64,) ,   lowerCamelCase=2 ,   lowerCamelCase=32 ,   lowerCamelCase="silu" ,   lowerCamelCase=True ,   )							->Dict:



         '''simple docstring'''
         super().__init__()
         __a				     =    layers_per_block

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

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

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

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

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

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

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

         __a				     =    False






      def  __UpperCamelCase							(  self ,   lowerCamelCase      )							->Union[str, Any]:



         '''simple docstring'''
         __a				     =    x
         __a				     =    self.conv_in(__a      )

         if self.training and self.gradient_checkpointing:

            def create_custom_forward(lowerCamelCase      ):
               def custom_forward(*lowerCamelCase      ):
                  return module(*__a      )

               return custom_forward

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

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

            # middle
            __a				     =    self.mid_block(__a      )

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

         return sample







class 							__SCREAMING_SNAKE_CASE       (							nn.Module     ):




      def __init__(  self ,   lowerCamelCase=3 ,   lowerCamelCase=3 ,   lowerCamelCase=("UpDecoderBlock2D",) ,   lowerCamelCase=(64,) ,   lowerCamelCase=2 ,   lowerCamelCase=32 ,   lowerCamelCase="silu" ,   lowerCamelCase="group" ,   )							->Union[str, Any]:



         '''simple docstring'''
         super().__init__()
         __a				     =    layers_per_block

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

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

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

         # mid
         __a				     =    UNetMidBlockaD(
             in_channels=block_out_channels[-1] ,   resnet_eps=1e-6 ,   resnet_act_fn=__a ,   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=__a ,   temb_channels=__a ,   )

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

            __a				     =    i == len(__a      ) - 1

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

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

         __a				     =    False






      def  __UpperCamelCase							(  self ,   lowerCamelCase ,   lowerCamelCase=None      )							->List[Any]:



         '''simple docstring'''
         __a				     =    z
         __a				     =    self.conv_in(__a      )

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

            def create_custom_forward(lowerCamelCase      ):
               def custom_forward(*lowerCamelCase      ):
                  return module(*__a      )

               return custom_forward

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

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

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

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

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

         return sample







class 							__SCREAMING_SNAKE_CASE       (							nn.Module     ):




      def __init__(  self ,   lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase=None ,   lowerCamelCase="random" ,   lowerCamelCase=False ,   lowerCamelCase=True      )							->Optional[int]:



         '''simple docstring'''
         super().__init__()
         __a				     =    n_e
         __a				     =    vq_embed_dim
         __a				     =    beta
         __a				     =    legacy

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

         __a				     =    remap
         if self.remap is not None:
            self.register_buffer('used' ,   torch.tensor(np.load(self.remap      )      )      )
            __a				     =    self.used.shape[0]
            __a				     =    unknown_index  # "random" or "extra" or integer
            if self.unknown_index == "extra":
               __a				     =    self.re_embed
               __a				     =    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:
            __a				     =    n_e

         __a				     =    sane_index_shape




      def  __UpperCamelCase							(  self ,   lowerCamelCase      )							->Tuple:



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




      def  __UpperCamelCase							(  self ,   lowerCamelCase      )							->Optional[Any]:



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




      def  __UpperCamelCase							(  self ,   lowerCamelCase      )							->Dict:



         '''simple docstring'''
         # reshape z -> (batch, height, width, channel) and flatten
         __a				     =    z.permute(0 ,   2 ,   3 ,   1      ).contiguous()
         __a				     =    z.view(-1 ,   self.vq_embed_dim      )

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

         __a				     =    self.embedding(__a      ).view(z.shape      )
         __a				     =    None
         __a				     =    None

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

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

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

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

         if self.sane_index_shape:
            __a				     =    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  __UpperCamelCase							(  self ,   lowerCamelCase ,   lowerCamelCase      )							->Union[str, Any]:



         '''simple docstring'''
         # shape specifying (batch, height, width, channel)
         if self.remap is not None:
            __a				     =    indices.reshape(shape[0] ,   -1      )  # add batch axis
            __a				     =    self.unmap_to_all(__a      )
            __a				     =    indices.reshape(-1      )  # flatten again

         # get quantized latent vectors
         __a				     =    self.embedding(__a      )

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

         return z_q







class 							__SCREAMING_SNAKE_CASE       (							__lowercase     ):




      def __init__(  self ,   lowerCamelCase ,   lowerCamelCase=False      )							->str:



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




      def  __UpperCamelCase							(  self ,   lowerCamelCase = None      )							->torch.FloatTensor:



         '''simple docstring'''
         # make sure sample is on the same device as the parameters and has same dtype
         __a				     =    randn_tensor(
             self.mean.shape ,   generator=__a ,   device=self.parameters.device ,   dtype=self.parameters.dtype      )
         __a				     =    self.mean + self.std * sample
         return x




      def  __UpperCamelCase							(  self ,   lowerCamelCase=None      )							->Optional[Any]:



         '''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  __UpperCamelCase							(  self ,   lowerCamelCase ,   lowerCamelCase=[1, 2, 3]      )							->str:



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






      def  __UpperCamelCase							(  self      )							->Dict:



         '''simple docstring'''
         return self.mean








'''simple docstring'''



from collections import OrderedDict
from typing import Mapping

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


__UpperCamelCase      :       Dict      			=      logging.get_logger(__name__)

__UpperCamelCase      :       Optional[int]      			=      {
    """facebook/xmod-base""": """https://huggingface.co/facebook/xmod-base/resolve/main/config.json""",
    """facebook/xmod-large-prenorm""": """https://huggingface.co/facebook/xmod-large-prenorm/resolve/main/config.json""",
    """facebook/xmod-base-13-125k""": """https://huggingface.co/facebook/xmod-base-13-125k/resolve/main/config.json""",
    """facebook/xmod-base-30-125k""": """https://huggingface.co/facebook/xmod-base-30-125k/resolve/main/config.json""",
    """facebook/xmod-base-30-195k""": """https://huggingface.co/facebook/xmod-base-30-195k/resolve/main/config.json""",
    """facebook/xmod-base-60-125k""": """https://huggingface.co/facebook/xmod-base-60-125k/resolve/main/config.json""",
    """facebook/xmod-base-60-265k""": """https://huggingface.co/facebook/xmod-base-60-265k/resolve/main/config.json""",
    """facebook/xmod-base-75-125k""": """https://huggingface.co/facebook/xmod-base-75-125k/resolve/main/config.json""",
    """facebook/xmod-base-75-269k""": """https://huggingface.co/facebook/xmod-base-75-269k/resolve/main/config.json""",
}
class 							__SCREAMING_SNAKE_CASE       (							_lowerCAmelCase     ):
      __a  					="xmod"




      def __init__(  self ,   lowerCamelCase=3_0522 ,   lowerCamelCase=768 ,   lowerCamelCase=12 ,   lowerCamelCase=12 ,   lowerCamelCase=3072 ,   lowerCamelCase="gelu" ,   lowerCamelCase=0.1 ,   lowerCamelCase=0.1 ,   lowerCamelCase=512 ,   lowerCamelCase=2 ,   lowerCamelCase=0.02 ,   lowerCamelCase=1e-12 ,   lowerCamelCase=1 ,   lowerCamelCase=0 ,   lowerCamelCase=2 ,   lowerCamelCase="absolute" ,   lowerCamelCase=True ,   lowerCamelCase=None ,   lowerCamelCase=False ,   lowerCamelCase=2 ,   lowerCamelCase=False ,   lowerCamelCase=True ,   lowerCamelCase=True ,   lowerCamelCase=("en_XX",) ,   lowerCamelCase=None ,   **lowerCamelCase ,   )							->Any:



         '''simple docstring'''
         super().__init__(pad_token_id=lowerCamelCase ,   bos_token_id=lowerCamelCase ,   eos_token_id=lowerCamelCase ,   **lowerCamelCase      )

         __a				     =    vocab_size
         __a				     =    hidden_size
         __a				     =    num_hidden_layers
         __a				     =    num_attention_heads
         __a				     =    hidden_act
         __a				     =    intermediate_size
         __a				     =    hidden_dropout_prob
         __a				     =    attention_probs_dropout_prob
         __a				     =    max_position_embeddings
         __a				     =    type_vocab_size
         __a				     =    initializer_range
         __a				     =    layer_norm_eps
         __a				     =    position_embedding_type
         __a				     =    use_cache
         __a				     =    classifier_dropout
         __a				     =    pre_norm
         __a				     =    adapter_reduction_factor
         __a				     =    adapter_layer_norm
         __a				     =    adapter_reuse_layer_norm
         __a				     =    ln_before_adapter
         __a				     =    list(lowerCamelCase      )
         __a				     =    default_language







class 							__SCREAMING_SNAKE_CASE       (							_lowerCAmelCase     ):




      @property
      def  __UpperCamelCase							(  self      )							->Mapping[str, Mapping[int, str]]:



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