"""simple docstring"""






import flax.linen as nn
import jax
import jax.numpy as jnp


class     UpperCAmelCase						(nn.Module  ):



    """simple docstring"""

    _UpperCAmelCase	:int
    _UpperCAmelCase	:jnp.dtype						    = jnp.floataa






    def _snake_case			(   self							):
       lowercase__:   List[Any]					=      nn.Conv(
           self.out_channels    ,			kernel_size=(3, 3)    ,			strides=(1, 1)    ,			padding=((1, 1), (1, 1))    ,			dtype=self.dtype    ,			)






    def __call__(   self    ,			_UpperCAmelCase							):
       lowercase__, lowercase__, lowercase__, lowercase__:   List[str]					=      hidden_states.shape
       lowercase__:   Optional[int]					=      jax.image.resize(
           _UpperCAmelCase    ,			shape=(batch, height * 2, width * 2, channels)    ,			method='''nearest'''    ,			)
       lowercase__:   Any					=      self.conv(_UpperCAmelCase							)
       return hidden_states



class     UpperCAmelCase						(nn.Module  ):



    """simple docstring"""

    _UpperCAmelCase	:int
    _UpperCAmelCase	:jnp.dtype						    = jnp.floataa






    def _snake_case			(   self							):
       lowercase__:   Optional[Any]					=      nn.Conv(
           self.out_channels    ,			kernel_size=(3, 3)    ,			strides=(2, 2)    ,			padding=((1, 1), (1, 1))    ,			dtype=self.dtype    ,			)






    def __call__(   self    ,			_UpperCAmelCase							):
       # pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
       # hidden_states = jnp.pad(hidden_states, pad_width=pad)
       lowercase__:   str					=      self.conv(_UpperCAmelCase							)
       return hidden_states



class     UpperCAmelCase						(nn.Module  ):



    """simple docstring"""

    _UpperCAmelCase	:int
    _UpperCAmelCase	:int						    = None
    _UpperCAmelCase	:float						    = 0.0
    _UpperCAmelCase	:bool						    = None
    _UpperCAmelCase	:jnp.dtype						    = jnp.floataa






    def _snake_case			(   self							):
       lowercase__:   List[Any]					=      self.in_channels if self.out_channels is None else self.out_channels

       lowercase__:   Tuple					=      nn.GroupNorm(num_groups=32    ,			epsilon=1e-5							)
       lowercase__:   Any					=      nn.Conv(
           _UpperCAmelCase    ,			kernel_size=(3, 3)    ,			strides=(1, 1)    ,			padding=((1, 1), (1, 1))    ,			dtype=self.dtype    ,			)

       lowercase__:   List[str]					=      nn.Dense(_UpperCAmelCase    ,			dtype=self.dtype							)

       lowercase__:   List[Any]					=      nn.GroupNorm(num_groups=32    ,			epsilon=1e-5							)
       lowercase__:   Union[str, Any]					=      nn.Dropout(self.dropout_prob							)
       lowercase__:   int					=      nn.Conv(
           _UpperCAmelCase    ,			kernel_size=(3, 3)    ,			strides=(1, 1)    ,			padding=((1, 1), (1, 1))    ,			dtype=self.dtype    ,			)

       lowercase__:   Union[str, Any]					=      self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

       lowercase__:   Union[str, Any]					=      None
       if use_nin_shortcut:
          lowercase__:   List[str]					=      nn.Conv(
              _UpperCAmelCase    ,			kernel_size=(1, 1)    ,			strides=(1, 1)    ,			padding='''VALID'''    ,			dtype=self.dtype    ,			)






    def __call__(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase=True							):
       lowercase__:   List[str]					=      hidden_states
       lowercase__:   int					=      self.norma(_UpperCAmelCase							)
       lowercase__:   Any					=      nn.swish(_UpperCAmelCase							)
       lowercase__:   Optional[int]					=      self.conva(_UpperCAmelCase							)

       lowercase__:   Union[str, Any]					=      self.time_emb_proj(nn.swish(_UpperCAmelCase							)							)
       lowercase__:   List[str]					=      jnp.expand_dims(jnp.expand_dims(_UpperCAmelCase    ,			1							)    ,			1							)
       lowercase__:   Dict					=      hidden_states + temb

       lowercase__:   Tuple					=      self.norma(_UpperCAmelCase							)
       lowercase__:   List[str]					=      nn.swish(_UpperCAmelCase							)
       lowercase__:   Optional[int]					=      self.dropout(_UpperCAmelCase    ,			_UpperCAmelCase							)
       lowercase__:   Optional[int]					=      self.conva(_UpperCAmelCase							)

       if self.conv_shortcut is not None:
          lowercase__:   Dict					=      self.conv_shortcut(_UpperCAmelCase							)

       return hidden_states + residual

"""simple docstring"""






import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ....tokenization_utils_fast import PreTrainedTokenizerFast
from ....utils import logging
from .tokenization_retribert import RetriBertTokenizer


__A			   = logging.get_logger(__name__)

__A			   = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}

__A			   = {
    "vocab_file": {
        "yjernite/retribert-base-uncased": (
            "https://huggingface.co/yjernite/retribert-base-uncased/resolve/main/vocab.txt"
        ),
    },
    "tokenizer_file": {
        "yjernite/retribert-base-uncased": (
            "https://huggingface.co/yjernite/retribert-base-uncased/resolve/main/tokenizer.json"
        ),
    },
}

__A			   = {
    "yjernite/retribert-base-uncased": 5_1_2,
}


__A			   = {
    "yjernite/retribert-base-uncased": {"do_lower_case": True},
}


class     UpperCAmelCase						(_UpperCAmelCase  ):



    """simple docstring"""

    _UpperCAmelCase	:Dict						    = VOCAB_FILES_NAMES
    _UpperCAmelCase	:str						    = PRETRAINED_VOCAB_FILES_MAP
    _UpperCAmelCase	:Optional[int]						    = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    _UpperCAmelCase	:Optional[int]						    = PRETRAINED_INIT_CONFIGURATION
    _UpperCAmelCase	:str						    = RetriBertTokenizer
    _UpperCAmelCase	:List[str]						    = ["input_ids", "attention_mask"]






    def __init__(   self    ,			_UpperCAmelCase=None    ,			_UpperCAmelCase=None    ,			_UpperCAmelCase=True    ,			_UpperCAmelCase="[UNK]"    ,			_UpperCAmelCase="[SEP]"    ,			_UpperCAmelCase="[PAD]"    ,			_UpperCAmelCase="[CLS]"    ,			_UpperCAmelCase="[MASK]"    ,			_UpperCAmelCase=True    ,			_UpperCAmelCase=None    ,			**_UpperCAmelCase    ,			):
       super().__init__(
           _UpperCAmelCase    ,			tokenizer_file=_UpperCAmelCase    ,			do_lower_case=_UpperCAmelCase    ,			unk_token=_UpperCAmelCase    ,			sep_token=_UpperCAmelCase    ,			pad_token=_UpperCAmelCase    ,			cls_token=_UpperCAmelCase    ,			mask_token=_UpperCAmelCase    ,			tokenize_chinese_chars=_UpperCAmelCase    ,			strip_accents=_UpperCAmelCase    ,			**_UpperCAmelCase    ,			)

       lowercase__:   Optional[Any]					=      json.loads(self.backend_tokenizer.normalizer.__getstate__()							)
       if (
           normalizer_state.get('''lowercase'''    ,			_UpperCAmelCase							) != do_lower_case
           or normalizer_state.get('''strip_accents'''    ,			_UpperCAmelCase							) != strip_accents
           or normalizer_state.get('''handle_chinese_chars'''    ,			_UpperCAmelCase							) != tokenize_chinese_chars
       ):
          lowercase__:   str					=      getattr(_UpperCAmelCase    ,			normalizer_state.pop('''type'''							)							)
          lowercase__:   Optional[Any]					=      do_lower_case
          lowercase__:   Tuple					=      strip_accents
          lowercase__:   str					=      tokenize_chinese_chars
          lowercase__:   Union[str, Any]					=      normalizer_class(**_UpperCAmelCase							)

       lowercase__:   List[str]					=      do_lower_case






    def _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase=None							):
       lowercase__:   Union[str, Any]					=      [self.cls_token_id] + token_ids_a + [self.sep_token_id]

       if token_ids_a:
          output += token_ids_a + [self.sep_token_id]

       return output






    def _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase = None							):
       lowercase__:   Optional[int]					=      [self.sep_token_id]
       lowercase__:   Union[str, 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							) * [0] + len(token_ids_a + sep							) * [1]






    def _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase = None							):
       lowercase__:   List[str]					=      self._tokenizer.model.save(_UpperCAmelCase    ,			name=_UpperCAmelCase							)
       return tuple(_UpperCAmelCase							)
