'''simple docstring'''


import functools
import math

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


def    _lowerCamelCase       ( lowerCamelCase_   :				List[Any]      ,	lowerCamelCase_   :				List[str]      ,	lowerCamelCase_   :				Optional[int]      ,	lowerCamelCase_   :				Dict      ,	lowerCamelCase_   :				int = 4096  ):



     """simple docstring"""





     UpperCAmelCase_  ,     UpperCAmelCase_  ,     UpperCAmelCase_				:  List[Any]  = key.shape[-3:]
     UpperCAmelCase_				:  List[Any]  = value.shape[-1]
     UpperCAmelCase_				:  Any  = min(lowerCamelCase_      ,	lowerCamelCase_  )
     UpperCAmelCase_				:  Optional[int]  = query / jnp.sqrt(lowerCamelCase_  )

     @functools.partial(jax.checkpoint      ,	prevent_cse=lowerCamelCase_  )
     def summarize_chunk(lowerCamelCase_   :				List[Any]      ,	lowerCamelCase_   :				int      ,	lowerCamelCase_   :				List[str]  ):
          UpperCAmelCase_				:  List[Any]  = jnp.einsum('...qhd,...khd->...qhk'      ,	lowerCamelCase_      ,	lowerCamelCase_      ,	precision=lowerCamelCase_  )

          UpperCAmelCase_				:  Optional[Any]  = jnp.max(lowerCamelCase_      ,	axis=-1      ,	keepdims=lowerCamelCase_  )
          UpperCAmelCase_				:  str  = jax.lax.stop_gradient(lowerCamelCase_  )
          UpperCAmelCase_				:  Optional[int]  = jnp.exp(attn_weights - max_score  )

          UpperCAmelCase_				:  List[Any]  = jnp.einsum('...vhf,...qhv->...qhf'      ,	lowerCamelCase_      ,	lowerCamelCase_      ,	precision=lowerCamelCase_  )
          UpperCAmelCase_				:  List[Any]  = jnp.einsum('...qhk->...qh'      ,	lowerCamelCase_  )

          return (exp_values, exp_weights.sum(axis=-1  ), max_score)

     def chunk_scanner(lowerCamelCase_   :				Union[str, Any]  ):
          # julienne key array
          UpperCAmelCase_				:  Tuple  = jax.lax.dynamic_slice(
              operand=lowerCamelCase_      ,	start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]      ,	slice_sizes=list(key.shape[:-3]  ) + [key_chunk_size, num_heads, k_features]      ,	)

          # julienne value array
          UpperCAmelCase_				:  Union[str, Any]  = jax.lax.dynamic_slice(
              operand=lowerCamelCase_      ,	start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]      ,	slice_sizes=list(value.shape[:-3]  ) + [key_chunk_size, num_heads, v_features]      ,	)

          return summarize_chunk(lowerCamelCase_      ,	lowerCamelCase_      ,	lowerCamelCase_  )

     UpperCAmelCase_  ,     UpperCAmelCase_  ,     UpperCAmelCase_				:  int  = jax.lax.map(f=lowerCamelCase_      ,	xs=jnp.arange(0      ,	lowerCamelCase_      ,	lowerCamelCase_  )  )

     UpperCAmelCase_				:  str  = jnp.max(lowerCamelCase_      ,	axis=0      ,	keepdims=lowerCamelCase_  )
     UpperCAmelCase_				:  List[Any]  = jnp.exp(chunk_max - global_max  )

     chunk_values *= jnp.expand_dims(lowerCamelCase_      ,	axis=-1  )
     chunk_weights *= max_diffs

     UpperCAmelCase_				:  Dict  = chunk_values.sum(axis=0  )
     UpperCAmelCase_				:  Tuple  = jnp.expand_dims(lowerCamelCase_      ,	-1  ).sum(axis=0  )

     return all_values / all_weights


def    _lowerCamelCase       ( lowerCamelCase_   :				List[str]      ,	lowerCamelCase_   :				Union[str, Any]      ,	lowerCamelCase_   :				Tuple      ,	lowerCamelCase_   :				Optional[Any]=jax.lax.Precision.HIGHEST      ,	lowerCamelCase_   :				int = 1024      ,	lowerCamelCase_   :				int = 4096  ):



     """simple docstring"""





     UpperCAmelCase_  ,     UpperCAmelCase_  ,     UpperCAmelCase_				:  Union[str, Any]  = query.shape[-3:]

     def chunk_scanner(lowerCamelCase_   :				Union[str, Any]      ,	lowerCamelCase_   :				List[str]  ):
          # julienne query array
          UpperCAmelCase_				:  Tuple  = jax.lax.dynamic_slice(
              operand=lowerCamelCase_      ,	start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]      ,	slice_sizes=list(query.shape[:-3]  ) + [min(lowerCamelCase_      ,	lowerCamelCase_  ), num_heads, q_features]      ,	)

          return (
              chunk_idx + query_chunk_size,  # unused ignore it
              _query_chunk_attention(
                  query=lowerCamelCase_      ,	key=lowerCamelCase_      ,	value=lowerCamelCase_      ,	precision=lowerCamelCase_      ,	key_chunk_size=lowerCamelCase_  ),
          )

     UpperCAmelCase_  ,     UpperCAmelCase_				:  Optional[int]  = jax.lax.scan(
         f=lowerCamelCase_      ,	init=0      ,	xs=lowerCamelCase_      ,	length=math.ceil(num_q / query_chunk_size  )  )

     return jnp.concatenate(lowerCamelCase_      ,	axis=-3  )  # fuse the chunked result back






class 	__SCREAMING_SNAKE_CASE		(			nn.Module			):







  '''simple docstring'''






  lowerCamelCase_						:int
  lowerCamelCase_						:int		       =						8
  lowerCamelCase_						:int		       =						64
  lowerCamelCase_						:float		       =						0.0
  lowerCamelCase_						:bool		       =						False
  lowerCamelCase_						:jnp.dtype		       =						jnp.floataa


  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       UpperCAmelCase_				:  List[str]  = self.dim_head * self.heads
       UpperCAmelCase_				:  Optional[Any]  = self.dim_head**-0.5

       # Weights were exported with old names {to_q, to_k, to_v, to_out}
       UpperCAmelCase_				:  int  = nn.Dense(snake_case_ ,						use_bias=snake_case_ ,						dtype=self.dtype ,						name='to_q'		)
       UpperCAmelCase_				:  Dict  = nn.Dense(snake_case_ ,						use_bias=snake_case_ ,						dtype=self.dtype ,						name='to_k'		)
       UpperCAmelCase_				:  Dict  = nn.Dense(snake_case_ ,						use_bias=snake_case_ ,						dtype=self.dtype ,						name='to_v'		)

       UpperCAmelCase_				:  Optional[Any]  = nn.Dense(self.query_dim ,						dtype=self.dtype ,						name='to_out_0'		)
       UpperCAmelCase_				:  Any  = nn.Dropout(rate=self.dropout		)


  def _UpperCamelCase	(					self ,						snake_case_		):




       '''simple docstring'''

       UpperCAmelCase_  ,     UpperCAmelCase_  ,     UpperCAmelCase_				:  List[str]  = tensor.shape
       UpperCAmelCase_				:  Union[str, Any]  = self.heads
       UpperCAmelCase_				:  Optional[int]  = tensor.reshape(snake_case_ ,						snake_case_ ,						snake_case_ ,						dim // head_size		)
       UpperCAmelCase_				:  Optional[int]  = jnp.transpose(snake_case_ ,						(0, 2, 1, 3)		)
       UpperCAmelCase_				:  Dict  = tensor.reshape(batch_size * head_size ,						snake_case_ ,						dim // head_size		)
       return tensor


  def _UpperCamelCase	(					self ,						snake_case_		):




       '''simple docstring'''

       UpperCAmelCase_  ,     UpperCAmelCase_  ,     UpperCAmelCase_				:  Dict  = tensor.shape
       UpperCAmelCase_				:  Any  = self.heads
       UpperCAmelCase_				:  List[str]  = tensor.reshape(batch_size // head_size ,						snake_case_ ,						snake_case_ ,						snake_case_		)
       UpperCAmelCase_				:  Tuple  = jnp.transpose(snake_case_ ,						(0, 2, 1, 3)		)
       UpperCAmelCase_				:  Union[str, Any]  = tensor.reshape(batch_size // head_size ,						snake_case_ ,						dim * head_size		)
       return tensor






  def __call__(					self ,						snake_case_ ,						snake_case_=None ,						snake_case_=True		):




       '''simple docstring'''

       UpperCAmelCase_				:  Optional[int]  = hidden_states if context is None else context

       UpperCAmelCase_				:  Tuple  = self.query(snake_case_		)
       UpperCAmelCase_				:  str  = self.key(snake_case_		)
       UpperCAmelCase_				:  Union[str, Any]  = self.value(snake_case_		)

       UpperCAmelCase_				:  str  = self.reshape_heads_to_batch_dim(snake_case_		)
       UpperCAmelCase_				:  List[str]  = self.reshape_heads_to_batch_dim(snake_case_		)
       UpperCAmelCase_				:  List[Any]  = self.reshape_heads_to_batch_dim(snake_case_		)

       if self.use_memory_efficient_attention:
            UpperCAmelCase_				:  int  = query_states.transpose(1 ,						0 ,						2		)
            UpperCAmelCase_				:  int  = key_states.transpose(1 ,						0 ,						2		)
            UpperCAmelCase_				:  Optional[int]  = value_states.transpose(1 ,						0 ,						2		)

            # this if statement create a chunk size for each layer of the unet
            # the chunk size is equal to the query_length dimension of the deepest layer of the unet

            UpperCAmelCase_				:  List[Any]  = query_states.shape[-3]
            if flatten_latent_dim % 6_4 == 0:
                 UpperCAmelCase_				:  Union[str, Any]  = int(flatten_latent_dim / 6_4		)
            elif flatten_latent_dim % 1_6 == 0:
                 UpperCAmelCase_				:  Union[str, Any]  = int(flatten_latent_dim / 1_6		)
            elif flatten_latent_dim % 4 == 0:
                 UpperCAmelCase_				:  Any  = int(flatten_latent_dim / 4		)
            else:
                 UpperCAmelCase_				:  List[str]  = int(snake_case_		)

            UpperCAmelCase_				:  Optional[int]  = jax_memory_efficient_attention(
                snake_case_ ,						snake_case_ ,						snake_case_ ,						query_chunk_size=snake_case_ ,						key_chunk_size=4_0_9_6 * 4		)

            UpperCAmelCase_				:  Optional[Any]  = hidden_states.transpose(1 ,						0 ,						2		)
       else:
            # compute attentions
            UpperCAmelCase_				:  Optional[int]  = jnp.einsum('b i d, b j d->b i j' ,						snake_case_ ,						snake_case_		)
            UpperCAmelCase_				:  Union[str, Any]  = attention_scores * self.scale
            UpperCAmelCase_				:  Dict  = nn.softmax(snake_case_ ,						axis=2		)

            # attend to values
            UpperCAmelCase_				:  Tuple  = jnp.einsum('b i j, b j d -> b i d' ,						snake_case_ ,						snake_case_		)

       UpperCAmelCase_				:  int  = self.reshape_batch_dim_to_heads(snake_case_		)
       UpperCAmelCase_				:  List[str]  = self.proj_attn(snake_case_		)
       return self.dropout_layer(snake_case_ ,						deterministic=snake_case_		)







class 	__SCREAMING_SNAKE_CASE		(			nn.Module			):







  '''simple docstring'''






  lowerCamelCase_						:int
  lowerCamelCase_						:int
  lowerCamelCase_						:int
  lowerCamelCase_						:float		       =						0.0
  lowerCamelCase_						:bool		       =						False
  lowerCamelCase_						:jnp.dtype		       =						jnp.floataa
  lowerCamelCase_						:bool		       =						False


  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       UpperCAmelCase_				:  Union[str, Any]  = FlaxAttention(
           self.dim ,						self.n_heads ,						self.d_head ,						self.dropout ,						self.use_memory_efficient_attention ,						dtype=self.dtype		)
       # cross attention
       UpperCAmelCase_				:  Any  = FlaxAttention(
           self.dim ,						self.n_heads ,						self.d_head ,						self.dropout ,						self.use_memory_efficient_attention ,						dtype=self.dtype		)
       UpperCAmelCase_				:  List[str]  = FlaxFeedForward(dim=self.dim ,						dropout=self.dropout ,						dtype=self.dtype		)
       UpperCAmelCase_				:  Tuple  = nn.LayerNorm(epsilon=1E-5 ,						dtype=self.dtype		)
       UpperCAmelCase_				:  int  = nn.LayerNorm(epsilon=1E-5 ,						dtype=self.dtype		)
       UpperCAmelCase_				:  Optional[int]  = nn.LayerNorm(epsilon=1E-5 ,						dtype=self.dtype		)
       UpperCAmelCase_				:  int  = nn.Dropout(rate=self.dropout		)






  def __call__(					self ,						snake_case_ ,						snake_case_ ,						snake_case_=True		):




       '''simple docstring'''

       UpperCAmelCase_				:  str  = hidden_states
       if self.only_cross_attention:
            UpperCAmelCase_				:  int  = self.attna(self.norma(snake_case_		) ,						snake_case_ ,						deterministic=snake_case_		)
       else:
            UpperCAmelCase_				:  Any  = self.attna(self.norma(snake_case_		) ,						deterministic=snake_case_		)
       UpperCAmelCase_				:  List[Any]  = hidden_states + residual

       # cross attention
       UpperCAmelCase_				:  List[str]  = hidden_states
       UpperCAmelCase_				:  Optional[Any]  = self.attna(self.norma(snake_case_		) ,						snake_case_ ,						deterministic=snake_case_		)
       UpperCAmelCase_				:  List[str]  = hidden_states + residual

       # feed forward
       UpperCAmelCase_				:  Optional[int]  = hidden_states
       UpperCAmelCase_				:  int  = self.ff(self.norma(snake_case_		) ,						deterministic=snake_case_		)
       UpperCAmelCase_				:  List[Any]  = hidden_states + residual

       return self.dropout_layer(snake_case_ ,						deterministic=snake_case_		)







class 	__SCREAMING_SNAKE_CASE		(			nn.Module			):







  '''simple docstring'''






  lowerCamelCase_						:int
  lowerCamelCase_						:int
  lowerCamelCase_						:int
  lowerCamelCase_						:int		       =						1
  lowerCamelCase_						:float		       =						0.0
  lowerCamelCase_						:bool		       =						False
  lowerCamelCase_						:bool		       =						False
  lowerCamelCase_						:jnp.dtype		       =						jnp.floataa
  lowerCamelCase_						:bool		       =						False


  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       UpperCAmelCase_				:  Tuple  = nn.GroupNorm(num_groups=3_2 ,						epsilon=1E-5		)

       UpperCAmelCase_				:  Any  = self.n_heads * self.d_head
       if self.use_linear_projection:
            UpperCAmelCase_				:  int  = nn.Dense(snake_case_ ,						dtype=self.dtype		)
       else:
            UpperCAmelCase_				:  Any  = nn.Conv(
                snake_case_ ,						kernel_size=(1, 1) ,						strides=(1, 1) ,						padding='VALID' ,						dtype=self.dtype ,						)

       UpperCAmelCase_				:  Dict  = [
           FlaxBasicTransformerBlock(
               snake_case_ ,						self.n_heads ,						self.d_head ,						dropout=self.dropout ,						only_cross_attention=self.only_cross_attention ,						dtype=self.dtype ,						use_memory_efficient_attention=self.use_memory_efficient_attention ,						)
           for _ in range(self.depth		)
       ]

       if self.use_linear_projection:
            UpperCAmelCase_				:  Union[str, Any]  = nn.Dense(snake_case_ ,						dtype=self.dtype		)
       else:
            UpperCAmelCase_				:  List[Any]  = nn.Conv(
                snake_case_ ,						kernel_size=(1, 1) ,						strides=(1, 1) ,						padding='VALID' ,						dtype=self.dtype ,						)

       UpperCAmelCase_				:  str  = nn.Dropout(rate=self.dropout		)






  def __call__(					self ,						snake_case_ ,						snake_case_ ,						snake_case_=True		):




       '''simple docstring'''

       UpperCAmelCase_  ,     UpperCAmelCase_  ,     UpperCAmelCase_  ,     UpperCAmelCase_				:  Union[str, Any]  = hidden_states.shape
       UpperCAmelCase_				:  Union[str, Any]  = hidden_states
       UpperCAmelCase_				:  Optional[Any]  = self.norm(snake_case_		)
       if self.use_linear_projection:
            UpperCAmelCase_				:  str  = hidden_states.reshape(snake_case_ ,						height * width ,						snake_case_		)
            UpperCAmelCase_				:  Tuple  = self.proj_in(snake_case_		)
       else:
            UpperCAmelCase_				:  List[str]  = self.proj_in(snake_case_		)
            UpperCAmelCase_				:  Optional[Any]  = hidden_states.reshape(snake_case_ ,						height * width ,						snake_case_		)

       for transformer_block in self.transformer_blocks:
            UpperCAmelCase_				:  Optional[int]  = transformer_block(snake_case_ ,						snake_case_ ,						deterministic=snake_case_		)

       if self.use_linear_projection:
            UpperCAmelCase_				:  List[Any]  = self.proj_out(snake_case_		)
            UpperCAmelCase_				:  Optional[int]  = hidden_states.reshape(snake_case_ ,						snake_case_ ,						snake_case_ ,						snake_case_		)
       else:
            UpperCAmelCase_				:  Optional[Any]  = hidden_states.reshape(snake_case_ ,						snake_case_ ,						snake_case_ ,						snake_case_		)
            UpperCAmelCase_				:  Dict  = self.proj_out(snake_case_		)

       UpperCAmelCase_				:  Optional[int]  = hidden_states + residual
       return self.dropout_layer(snake_case_ ,						deterministic=snake_case_		)







class 	__SCREAMING_SNAKE_CASE		(			nn.Module			):







  '''simple docstring'''






  lowerCamelCase_						:int
  lowerCamelCase_						:float		       =						0.0
  lowerCamelCase_						:jnp.dtype		       =						jnp.floataa


  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       UpperCAmelCase_				:  Union[str, Any]  = FlaxGEGLU(self.dim ,						self.dropout ,						self.dtype		)
       UpperCAmelCase_				:  List[Any]  = nn.Dense(self.dim ,						dtype=self.dtype		)






  def __call__(					self ,						snake_case_ ,						snake_case_=True		):




       '''simple docstring'''

       UpperCAmelCase_				:  str  = self.net_a(snake_case_ ,						deterministic=snake_case_		)
       UpperCAmelCase_				:  Any  = self.net_a(snake_case_		)
       return hidden_states







class 	__SCREAMING_SNAKE_CASE		(			nn.Module			):







  '''simple docstring'''






  lowerCamelCase_						:int
  lowerCamelCase_						:float		       =						0.0
  lowerCamelCase_						:jnp.dtype		       =						jnp.floataa


  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       UpperCAmelCase_				:  Tuple  = self.dim * 4
       UpperCAmelCase_				:  Any  = nn.Dense(inner_dim * 2 ,						dtype=self.dtype		)
       UpperCAmelCase_				:  Any  = nn.Dropout(rate=self.dropout		)






  def __call__(					self ,						snake_case_ ,						snake_case_=True		):




       '''simple docstring'''

       UpperCAmelCase_				:  Dict  = self.proj(snake_case_		)
       UpperCAmelCase_  ,     UpperCAmelCase_				:  Any  = jnp.split(snake_case_ ,						2 ,						axis=2		)
       return self.dropout_layer(hidden_linear * nn.gelu(snake_case_		) ,						deterministic=snake_case_		)












import os
from typing import List, Union

import tensorflow as tf
from tensorflow_text import BertTokenizer as BertTokenizerLayer
from tensorflow_text import FastBertTokenizer, ShrinkLongestTrimmer, case_fold_utfa, combine_segments, pad_model_inputs

from .tokenization_bert import BertTokenizer
class       _UpperCamelCase   (					tf.keras.layers.Layer		):
      def __init__(      self :int			,   lowerCamelCase :List			,   lowerCamelCase :bool			,   lowerCamelCase :int = None			,   lowerCamelCase :int = None			,   lowerCamelCase :int = None			,   lowerCamelCase :str = "longest"			,   lowerCamelCase :bool = True			,   lowerCamelCase :int = 512			,   lowerCamelCase :int = None			,   lowerCamelCase :bool = True			,   lowerCamelCase :bool = True			,   lowerCamelCase :bool = True			,   **lowerCamelCase :Optional[Any]			,   )  ->      List[str]:
       super().__init__()
       if use_fast_bert_tokenizer:
        UpperCAmelCase__							=    FastBertTokenizer(
            lowerCamelCase			,   token_out_type=tf.intaa			,   lower_case_nfd_strip_accents=lowerCamelCase			,   **lowerCamelCase   )
       else:
        UpperCAmelCase__							=    tf.lookup.StaticVocabularyTable(
            tf.lookup.KeyValueTensorInitializer(
                keys=lowerCamelCase			,   key_dtype=tf.string			,   values=tf.range(tf.size(lowerCamelCase			,   out_type=tf.intaa   )			,   dtype=tf.intaa   )			,   value_dtype=tf.intaa			,   )			,   num_oov_buckets=1			,   )
        UpperCAmelCase__							=    BertTokenizerLayer(
            lowerCamelCase			,   token_out_type=tf.intaa			,   lower_case=lowerCamelCase			,   **lowerCamelCase   )

       UpperCAmelCase__							=    vocab_list
       UpperCAmelCase__							=    do_lower_case
       UpperCAmelCase__							=    cls_token_id or vocab_list.index("[CLS]"   )
       UpperCAmelCase__							=    sep_token_id or vocab_list.index("[SEP]"   )
       UpperCAmelCase__							=    pad_token_id or vocab_list.index("[PAD]"   )
       UpperCAmelCase__							=    ShrinkLongestTrimmer(max_length - 3			,   axis=1   )  # Allow room for special tokens
       UpperCAmelCase__							=    max_length
       UpperCAmelCase__							=    padding
       UpperCAmelCase__							=    truncation
       UpperCAmelCase__							=    pad_to_multiple_of
       UpperCAmelCase__							=    return_token_type_ids
       UpperCAmelCase__							=    return_attention_mask
      @classmethod
      def 						UpperCAmelCase_    (      cls :List[str]			,   lowerCamelCase :"PreTrainedTokenizerBase"			,   **lowerCamelCase :Dict   )  ->      Dict:  # noqa: F821
       UpperCAmelCase__							=    kwargs.pop("do_lower_case"			,   lowerCamelCase   )
       UpperCAmelCase__							=    tokenizer.do_lower_case if do_lower_case is None else do_lower_case
       UpperCAmelCase__							=    kwargs.pop("cls_token_id"			,   lowerCamelCase   )
       UpperCAmelCase__							=    tokenizer.cls_token_id if cls_token_id is None else cls_token_id
       UpperCAmelCase__							=    kwargs.pop("sep_token_id"			,   lowerCamelCase   )
       UpperCAmelCase__							=    tokenizer.sep_token_id if sep_token_id is None else sep_token_id
       UpperCAmelCase__							=    kwargs.pop("pad_token_id"			,   lowerCamelCase   )
       UpperCAmelCase__							=    tokenizer.pad_token_id if pad_token_id is None else pad_token_id

       UpperCAmelCase__							=    tokenizer.get_vocab()
       UpperCAmelCase__							=    sorted([(wordpiece, idx) for wordpiece, idx in vocab.items()]			,   key=lambda lowerCamelCase   : x[1]   )
       UpperCAmelCase__							=    [entry[0] for entry in vocab]
       return cls(
           vocab_list=lowerCamelCase			,   do_lower_case=lowerCamelCase			,   cls_token_id=lowerCamelCase			,   sep_token_id=lowerCamelCase			,   pad_token_id=lowerCamelCase			,   **lowerCamelCase			,   )
      @classmethod
      def 						UpperCAmelCase_    (      cls :Any			,   lowerCamelCase :Union[str, os.PathLike]			,   *lowerCamelCase :str			,   **lowerCamelCase :List[str]   )  ->      Tuple:
       try:
        UpperCAmelCase__							=    BertTokenizer.from_pretrained(lowerCamelCase			,   *lowerCamelCase			,   **lowerCamelCase   )
       except:  # noqa: E722
        from .tokenization_bert_fast import BertTokenizerFast

        UpperCAmelCase__							=    BertTokenizerFast.from_pretrained(lowerCamelCase			,   *lowerCamelCase			,   **lowerCamelCase   )
       return cls.from_tokenizer(lowerCamelCase			,   **lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :int			,   lowerCamelCase :Optional[int]   )  ->      List[str]:
       if self.do_lower_case:
        UpperCAmelCase__							=    case_fold_utfa(lowerCamelCase   )
       UpperCAmelCase__							=    self.tf_tokenizer.tokenize(lowerCamelCase   )
       return tokens.merge_dims(1			,   -1   )
      def 						UpperCAmelCase_    (      self :Optional[Any]			,   lowerCamelCase :Tuple			,   lowerCamelCase :str=None			,   lowerCamelCase :Tuple=None			,   lowerCamelCase :Any=None			,   lowerCamelCase :Optional[int]=None			,   lowerCamelCase :Optional[Any]=None			,   lowerCamelCase :Optional[Any]=None			,   lowerCamelCase :Tuple=None			,   )  ->      str:
       if padding is None:
        UpperCAmelCase__							=    self.padding
       if padding not in ("longest", "max_length"):
        raise ValueError("Padding must be either 'longest' or 'max_length'!"   )
       if max_length is not None and text_pair is not None:
        # Because we have to instantiate a Trimmer to do it properly
        raise ValueError("max_length cannot be overridden at call time when truncating paired texts!"   )
       if max_length is None:
        UpperCAmelCase__							=    self.max_length
       if truncation is None:
        UpperCAmelCase__							=    self.truncation
       if pad_to_multiple_of is None:
        UpperCAmelCase__							=    self.pad_to_multiple_of
       if return_token_type_ids is None:
        UpperCAmelCase__							=    self.return_token_type_ids
       if return_attention_mask is None:
        UpperCAmelCase__							=    self.return_attention_mask
       if not isinstance(lowerCamelCase			,   tf.Tensor   ):
        UpperCAmelCase__							=    tf.convert_to_tensor(lowerCamelCase   )
       if text_pair is not None and not isinstance(lowerCamelCase			,   tf.Tensor   ):
        UpperCAmelCase__							=    tf.convert_to_tensor(lowerCamelCase   )
       if text_pair is not None:
        if text.shape.rank > 1:
         raise ValueError("text argument should not be multidimensional when a text pair is supplied!"   )
        if text_pair.shape.rank > 1:
         raise ValueError("text_pair should not be multidimensional!"   )
       if text.shape.rank == 2:
        UpperCAmelCase__			,			UpperCAmelCase__							=    text[:, 0], text[:, 1]
       UpperCAmelCase__							=    self.unpaired_tokenize(lowerCamelCase   )
       if text_pair is None:  # Unpaired text
        if truncation:
         UpperCAmelCase__							=    text[:, : max_length - 2]  # Allow room for special tokens
        UpperCAmelCase__			,			UpperCAmelCase__							=    combine_segments(
            (text,)			,   start_of_sequence_id=self.cls_token_id			,   end_of_segment_id=self.sep_token_id   )
       else:  # Paired text
        UpperCAmelCase__							=    self.unpaired_tokenize(lowerCamelCase   )
        if truncation:
         UpperCAmelCase__			,			UpperCAmelCase__							=    self.paired_trimmer.trim([text, text_pair]   )
        UpperCAmelCase__			,			UpperCAmelCase__							=    combine_segments(
            (text, text_pair)			,   start_of_sequence_id=self.cls_token_id			,   end_of_segment_id=self.sep_token_id   )
       if padding == "longest":
        UpperCAmelCase__							=    input_ids.bounding_shape(axis=1   )
        if pad_to_multiple_of is not None:
         # No ceiling division in tensorflow, so we negate floordiv instead
         UpperCAmelCase__							=    pad_to_multiple_of * (-tf.math.floordiv(-pad_length			,   lowerCamelCase   ))
       else:
        UpperCAmelCase__							=    max_length

       UpperCAmelCase__			,			UpperCAmelCase__							=    pad_model_inputs(lowerCamelCase			,   max_seq_length=lowerCamelCase			,   pad_value=self.pad_token_id   )
       UpperCAmelCase__							=    {"input_ids": input_ids}
       if return_attention_mask:
        UpperCAmelCase__							=    attention_mask
       if return_token_type_ids:
        UpperCAmelCase__			,			UpperCAmelCase__							=    pad_model_inputs(
            lowerCamelCase			,   max_seq_length=lowerCamelCase			,   pad_value=self.pad_token_id   )
        UpperCAmelCase__							=    token_type_ids
       return output







      def 						UpperCAmelCase_    (      self :Tuple   )  ->      Tuple:
       return {
           "vocab_list": self.vocab_list,
           "do_lower_case": self.do_lower_case,
           "cls_token_id": self.cls_token_id,
           "sep_token_id": self.sep_token_id,
           "pad_token_id": self.pad_token_id,
       }





