





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




class      UpperCamelCase (     nn.Module   ):
    __UpperCamelCase         =42
    __UpperCamelCase         =jnp.floataa

    def UpperCamelCase			(			self :      Optional[int]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			nn.Conv(
         self.out_channels							,    kernel_size=(3, 3)							,    strides=(1, 1)							,    padding=((1, 1), (1, 1))							,    dtype=self.dtype							,    )




    def __call__(			self :      Tuple							,    snake_case__ :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			hidden_states.shape
     SCREAMING_SNAKE_CASE								=			jax.image.resize(
         __lowerCamelCase							,    shape=(batch, height * 2, width * 2, channels)							,    method='nearest'							,    )
     SCREAMING_SNAKE_CASE								=			self.conv(__lowerCamelCase    )
     return hidden_states




class      UpperCamelCase (     nn.Module   ):
    __UpperCamelCase         =42
    __UpperCamelCase         =jnp.floataa

    def UpperCamelCase			(			self :      List[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			nn.Conv(
         self.out_channels							,    kernel_size=(3, 3)							,    strides=(2, 2)							,    padding=((1, 1), (1, 1))							,    dtype=self.dtype							,    )




    def __call__(			self :      int							,    snake_case__ :      Tuple    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.conv(__lowerCamelCase    )
     return hidden_states




class      UpperCamelCase (     nn.Module   ):
    __UpperCamelCase         =42
    __UpperCamelCase         =None
    __UpperCamelCase         =0.0
    __UpperCamelCase         =None
    __UpperCamelCase         =jnp.floataa

    def UpperCamelCase			(			self :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.in_channels if self.out_channels is None else self.out_channels

     SCREAMING_SNAKE_CASE								=			nn.GroupNorm(num_groups=3_2							,    epsilon=1E-5    )
     SCREAMING_SNAKE_CASE								=			nn.Conv(
         __lowerCamelCase							,    kernel_size=(3, 3)							,    strides=(1, 1)							,    padding=((1, 1), (1, 1))							,    dtype=self.dtype							,    )

     SCREAMING_SNAKE_CASE								=			nn.Dense(__lowerCamelCase							,    dtype=self.dtype    )

     SCREAMING_SNAKE_CASE								=			nn.GroupNorm(num_groups=3_2							,    epsilon=1E-5    )
     SCREAMING_SNAKE_CASE								=			nn.Dropout(self.dropout_prob    )
     SCREAMING_SNAKE_CASE								=			nn.Conv(
         __lowerCamelCase							,    kernel_size=(3, 3)							,    strides=(1, 1)							,    padding=((1, 1), (1, 1))							,    dtype=self.dtype							,    )

     SCREAMING_SNAKE_CASE								=			self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

     SCREAMING_SNAKE_CASE								=			None
     if use_nin_shortcut:
      SCREAMING_SNAKE_CASE								=			nn.Conv(
          __lowerCamelCase							,    kernel_size=(1, 1)							,    strides=(1, 1)							,    padding='VALID'							,    dtype=self.dtype							,    )




    def __call__(			self :      Optional[int]							,    snake_case__ :      List[str]							,    snake_case__ :      Any							,    snake_case__ :      Optional[Any]=True    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			hidden_states
     SCREAMING_SNAKE_CASE								=			self.norma(__lowerCamelCase    )
     SCREAMING_SNAKE_CASE								=			nn.swish(__lowerCamelCase    )
     SCREAMING_SNAKE_CASE								=			self.conva(__lowerCamelCase    )

     SCREAMING_SNAKE_CASE								=			self.time_emb_proj(nn.swish(__lowerCamelCase    )    )
     SCREAMING_SNAKE_CASE								=			jnp.expand_dims(jnp.expand_dims(__lowerCamelCase							,    1    )							,    1    )
     SCREAMING_SNAKE_CASE								=			hidden_states + temb

     SCREAMING_SNAKE_CASE								=			self.norma(__lowerCamelCase    )
     SCREAMING_SNAKE_CASE								=			nn.swish(__lowerCamelCase    )
     SCREAMING_SNAKE_CASE								=			self.dropout(__lowerCamelCase							,    __lowerCamelCase    )
     SCREAMING_SNAKE_CASE								=			self.conva(__lowerCamelCase    )

     if self.conv_shortcut is not None:
      SCREAMING_SNAKE_CASE								=			self.conv_shortcut(__lowerCamelCase    )

     return hidden_states + residual




from typing import List, Optional, Tuple

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_herbert import HerbertTokenizer


a_    :						Optional[Any]    		=      logging.get_logger(__name__)

a_    :						Optional[Any]    		=      {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}

a_    :						Any    		=      {
    "vocab_file": {
        "allegro/herbert-base-cased": "https://huggingface.co/allegro/herbert-base-cased/resolve/main/vocab.json"
    },
    "merges_file": {
        "allegro/herbert-base-cased": "https://huggingface.co/allegro/herbert-base-cased/resolve/main/merges.txt"
    },
}

a_    :						Union[str, Any]    		=      {"allegro/herbert-base-cased": 514}
a_    :						List[Any]    		=      {}




class      UpperCamelCase (     SCREAMING_SNAKE_CASE   ):
    __UpperCamelCase         =VOCAB_FILES_NAMES
    __UpperCamelCase         =PRETRAINED_VOCAB_FILES_MAP
    __UpperCamelCase         =PRETRAINED_INIT_CONFIGURATION
    __UpperCamelCase         =PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    __UpperCamelCase         =HerbertTokenizer

    def __init__(			self :      Tuple							,    snake_case__ :      Optional[Any]=None							,    snake_case__ :      int=None							,    snake_case__ :      Optional[int]=None							,    snake_case__ :      str="<s>"							,    snake_case__ :      Tuple="<unk>"							,    snake_case__ :      List[str]="<pad>"							,    snake_case__ :      Tuple="<mask>"							,    snake_case__ :      Dict="</s>"							,    **snake_case__ :      List[str]							,    ):

     """simple docstring"""






     super().__init__(
         snake_case__							,    snake_case__							,    tokenizer_file=snake_case__							,    cls_token=snake_case__							,    unk_token=snake_case__							,    pad_token=snake_case__							,    mask_token=snake_case__							,    sep_token=snake_case__							,    **snake_case__							,    )

    def UpperCamelCase			(			self :      Union[str, Any]							,    snake_case__ :      List[int]							,    snake_case__ :      Optional[List[int]] = None    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			[self.cls_token_id]
     SCREAMING_SNAKE_CASE								=			[self.sep_token_id]
     if token_ids_a is None:
      return cls + token_ids_a + sep

     return cls + token_ids_a + sep + token_ids_a + sep

    def UpperCamelCase			(			self :      Optional[Any]							,    snake_case__ :      List[int]							,    snake_case__ :      Optional[List[int]] = None							,    snake_case__ :      bool = False    ):

     """simple docstring"""






     if already_has_special_tokens:
      return super().get_special_tokens_mask(
          token_ids_a=snake_case__							,    token_ids_a=snake_case__							,    already_has_special_tokens=snake_case__    )

     if token_ids_a is None:
      return [1] + ([0] * len(snake_case__    )) + [1]
     return [1] + ([0] * len(snake_case__    )) + [1] + ([0] * len(snake_case__    )) + [1]

    def UpperCamelCase			(			self :      Optional[Any]							,    snake_case__ :      List[int]							,    snake_case__ :      Optional[List[int]] = None    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			[self.sep_token_id]
     SCREAMING_SNAKE_CASE								=			[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 UpperCamelCase			(			self :      Tuple							,    snake_case__ :      str							,    snake_case__ :      Optional[str] = None    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self._tokenizer.model.save(snake_case__							,    name=snake_case__    )
     return tuple(snake_case__    )

