
'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


SCREAMING_SNAKE_CASE     =       {
    """configuration_rembert""": ["""REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """RemBertConfig""", """RemBertOnnxConfig"""]
}

try:
      if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE     =       ["""RemBertTokenizer"""]

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE     =       ["""RemBertTokenizerFast"""]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE     =       [
          """REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """RemBertForCausalLM""",
          """RemBertForMaskedLM""",
          """RemBertForMultipleChoice""",
          """RemBertForQuestionAnswering""",
          """RemBertForSequenceClassification""",
          """RemBertForTokenClassification""",
          """RemBertLayer""",
          """RemBertModel""",
          """RemBertPreTrainedModel""",
          """load_tf_weights_in_rembert""",
      ]


try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE     =       [
          """TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """TFRemBertForCausalLM""",
          """TFRemBertForMaskedLM""",
          """TFRemBertForMultipleChoice""",
          """TFRemBertForQuestionAnswering""",
          """TFRemBertForSequenceClassification""",
          """TFRemBertForTokenClassification""",
          """TFRemBertLayer""",
          """TFRemBertModel""",
          """TFRemBertPreTrainedModel""",
      ]


if TYPE_CHECKING:
      from .configuration_rembert import REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RemBertConfig, RemBertOnnxConfig

      try:
            if not is_sentencepiece_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_rembert import RemBertTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_rembert_fast import RemBertTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_rembert import (
                REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                RemBertForCausalLM,
                RemBertForMaskedLM,
                RemBertForMultipleChoice,
                RemBertForQuestionAnswering,
                RemBertForSequenceClassification,
                RemBertForTokenClassification,
                RemBertLayer,
                RemBertModel,
                RemBertPreTrainedModel,
                load_tf_weights_in_rembert,
            )

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_rembert import (
                TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                TFRemBertForCausalLM,
                TFRemBertForMaskedLM,
                TFRemBertForMultipleChoice,
                TFRemBertForQuestionAnswering,
                TFRemBertForSequenceClassification,
                TFRemBertForTokenClassification,
                TFRemBertLayer,
                TFRemBertModel,
                TFRemBertPreTrainedModel,
            )


else:
      import sys

      SCREAMING_SNAKE_CASE     =       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





'''simple docstring'''





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

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, Encoder, VectorQuantizer






@dataclass
class 							UpperCAmelCase_							(    A				):

 '''simple docstring'''
 lowercase_					:    torch.FloatTensor






class 							UpperCAmelCase_							(    A					,				A				):

 '''simple docstring'''







 @register_to_config
 def __init__(					self		:  List[str]			,				snake_case__		:  int = 3			,				snake_case__		:  int = 3			,				snake_case__		:  Tuple[str] = ("DownEncoderBlock2D",)			,				snake_case__		:  Tuple[str] = ("UpDecoderBlock2D",)			,				snake_case__		:  Tuple[int] = (64,)			,				snake_case__		:  int = 1			,				snake_case__		:  str = "silu"			,				snake_case__		:  int = 3			,				snake_case__		:  int = 32			,				snake_case__		:  int = 2_56			,				snake_case__		:  int = 32			,				snake_case__		:  Optional[int] = None			,				snake_case__		:  float = 0.18215			,				snake_case__		:  str = "group"			,				):



   '''simple docstring'''




   super().__init__()

   # pass init params to Encoder
   UpperCAmelCase__      :					Tuple				       =							Encoder(
       in_channels=snake_case__			,				out_channels=snake_case__			,				down_block_types=snake_case__			,				block_out_channels=snake_case__			,				layers_per_block=snake_case__			,				act_fn=snake_case__			,				norm_num_groups=snake_case__			,				double_z=snake_case__			,				)

   UpperCAmelCase__      :					str				       =							vq_embed_dim if vq_embed_dim is not None else latent_channels

   UpperCAmelCase__      :					Tuple				       =							nn.Convad(snake_case__			,				snake_case__			,				1 )
   UpperCAmelCase__      :					int				       =							VectorQuantizer(snake_case__			,				snake_case__			,				beta=0.25			,				remap=snake_case__			,				sane_index_shape=snake_case__ )
   UpperCAmelCase__      :					Tuple				       =							nn.Convad(snake_case__			,				snake_case__			,				1 )

   # pass init params to Decoder
   UpperCAmelCase__      :					Union[str, Any]				       =							Decoder(
       in_channels=snake_case__			,				out_channels=snake_case__			,				up_block_types=snake_case__			,				block_out_channels=snake_case__			,				layers_per_block=snake_case__			,				act_fn=snake_case__			,				norm_num_groups=snake_case__			,				norm_type=snake_case__			,				)







 @apply_forward_hook
 def 							UpperCamelCase  (					self		:  Union[str, Any]			,				snake_case__		:  torch.FloatTensor			,				snake_case__		:  bool = True ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[Any]				       =							self.encoder(snake_case__ )
   UpperCAmelCase__      :					List[str]				       =							self.quant_conv(snake_case__ )

   if not return_dict:
     return (h,)

   return VQEncoderOutput(latents=snake_case__ )







 @apply_forward_hook
 def 							UpperCamelCase  (					self		:  List[str]			,				snake_case__		:  torch.FloatTensor			,				snake_case__		:  bool = False			,				snake_case__		:  bool = True ):



   '''simple docstring'''




   if not force_not_quantize:
     UpperCAmelCase__       ,			UpperCAmelCase__       ,			UpperCAmelCase__      :					Any				       =							self.quantize(snake_case__ )
   else:
     UpperCAmelCase__      :					Tuple				       =							h
   UpperCAmelCase__      :					int				       =							self.post_quant_conv(snake_case__ )
   UpperCAmelCase__      :					Optional[int]				       =							self.decoder(snake_case__			,				quant if self.config.norm_type == "spatial" else None )

   if not return_dict:
     return (dec,)

   return DecoderOutput(sample=snake_case__ )







 def 							UpperCamelCase  (					self		:  Any			,				snake_case__		:  torch.FloatTensor			,				snake_case__		:  bool = True ):



   '''simple docstring'''




   UpperCAmelCase__      :					Dict				       =							sample
   UpperCAmelCase__      :					int				       =							self.encode(snake_case__ ).latents
   UpperCAmelCase__      :					Any				       =							self.decode(snake_case__ ).sample

   if not return_dict:
     return (dec,)

   return DecoderOutput(sample=snake_case__ )



