






"""simple docstring"""







import os
import re
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


__a       :	List[Any]             = logging.get_logger(__name__)

__a       :	Optional[int]             = {'vocab_file': 'spiece.model'}

__a       :	List[str]             = {
    'vocab_file': {
        'google/bigbird-roberta-base': 'https://huggingface.co/google/bigbird-roberta-base/resolve/main/spiece.model',
        'google/bigbird-roberta-large': (
            'https://huggingface.co/google/bigbird-roberta-large/resolve/main/spiece.model'
        ),
        'google/bigbird-base-trivia-itc': (
            'https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/spiece.model'
        ),
    }
}

__a       :	List[str]             = {
    'google/bigbird-roberta-base': 4096,
    'google/bigbird-roberta-large': 4096,
    'google/bigbird-base-trivia-itc': 4096,
}







class        _SCREAMING_SNAKE_CASE						(  __snake_case							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =VOCAB_FILES_NAMES
							_SCREAMING_SNAKE_CASE    =PRETRAINED_VOCAB_FILES_MAP
							_SCREAMING_SNAKE_CASE    =PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							_SCREAMING_SNAKE_CASE    =['input_ids', 'attention_mask']
							_SCREAMING_SNAKE_CASE    =[]




							def __init__(   self:     List[str]				,		__A:     Tuple				,		__A:     int="<unk>"				,		__A:     List[str]="<s>"				,		__A:     Any="</s>"				,		__A:     Union[str, Any]="<pad>"				,		__A:     Optional[Any]="[SEP]"				,		__A:     int="[MASK]"				,		__A:     int="[CLS]"				,		__A:     Optional[Dict[str, Any]] = None				,		**__A:     int				,		):

												'''simple docstring'''

												a__          =  AddedToken(__A				,		lstrip=__A				,		rstrip=__A    ) if isinstance(__A				,		__A    ) else bos_token
												a__          =  AddedToken(__A				,		lstrip=__A				,		rstrip=__A    ) if isinstance(__A				,		__A    ) else eos_token
												a__          =  AddedToken(__A				,		lstrip=__A				,		rstrip=__A    ) if isinstance(__A				,		__A    ) else unk_token
												a__          =  AddedToken(__A				,		lstrip=__A				,		rstrip=__A    ) if isinstance(__A				,		__A    ) else pad_token
												a__          =  AddedToken(__A				,		lstrip=__A				,		rstrip=__A    ) if isinstance(__A				,		__A    ) else cls_token
												a__          =  AddedToken(__A				,		lstrip=__A				,		rstrip=__A    ) if isinstance(__A				,		__A    ) else sep_token

												# Mask token behave like a normal word, i.e. include the space before it
												a__          =  AddedToken(__A				,		lstrip=__A				,		rstrip=__A    ) if isinstance(__A				,		__A    ) else mask_token

												a__          =  {} if sp_model_kwargs is None else sp_model_kwargs

												super().__init__(
												    bos_token=__A				,		eos_token=__A				,		unk_token=__A				,		pad_token=__A				,		sep_token=__A				,		mask_token=__A				,		cls_token=__A				,		sp_model_kwargs=self.sp_model_kwargs				,		**__A				,		)

												a__          =  vocab_file

												a__          =  spm.SentencePieceProcessor(**self.sp_model_kwargs    )
												self.sp_model.Load(__A    )




							@property
							def  lowercase		(   self:     Any    ):

												'''simple docstring'''

												return self.sp_model.get_piece_size()




							def  lowercase		(   self:     Union[str, Any]    ):

												'''simple docstring'''

												a__          =  {self.convert_ids_to_tokens(__A    ): i for i in range(self.vocab_size    )}
												vocab.update(self.added_tokens_encoder    )
												return vocab




							def __getstate__(   self:     List[Any]    ):

												'''simple docstring'''

												a__          =  self.__dict__.copy()
												a__          =  None
												return state




							def __setstate__(   self:     Optional[Any]				,		__A:     Tuple    ):

												'''simple docstring'''

												a__          =  d

												# for backward compatibility
												if not hasattr(self				,		'''sp_model_kwargs'''    ):
																	a__          =  {}

												a__          =  spm.SentencePieceProcessor(**self.sp_model_kwargs    )
												self.sp_model.Load(self.vocab_file    )




							def  lowercase		(   self:     Any				,		__A:     str    ):

												'''simple docstring'''

												return self.sp_model.encode(__A				,		out_type=__A    )




							def  lowercase		(   self:     List[str]				,		__A:     List[Any]    ):

												'''simple docstring'''

												return self.sp_model.piece_to_id(__A    )




							def  lowercase		(   self:     int				,		__A:     Dict    ):

												'''simple docstring'''

												a__          =  self.sp_model.IdToPiece(__A    )
												return token




							def  lowercase		(   self:     Optional[int]				,		__A:     int    ):

												'''simple docstring'''

												a__          =  []
												a__          =  ''''''
												a__          =  False
												for token in tokens:
																	# make sure that special tokens are not decoded using sentencepiece model
																	if token in self.all_special_tokens:
																						if not prev_is_special:
																											out_string += " "
																						out_string += self.sp_model.decode(__A    ) + token
																						a__          =  True
																						a__          =  []
																	else:
																						current_sub_tokens.append(__A    )
																						a__          =  False
												out_string += self.sp_model.decode(__A    )
												return out_string.strip()




							def  lowercase		(   self:     Tuple				,		__A:     List[int]				,		__A:     bool = False				,		__A:     bool = None				,		__A:     bool = True				,		**__A:     str				,		):

												'''simple docstring'''

												a__          =  kwargs.pop('''use_source_tokenizer'''				,		__A    )

												a__          =  self.convert_ids_to_tokens(__A				,		skip_special_tokens=__A    )

												# To avoid mixing byte-level and unicode for byte-level BPT
												# we need to build string separately for added tokens and byte-level tokens
												# cf. https://github.com/huggingface/transformers/issues/1133
												a__          =  []
												a__          =  []
												for token in filtered_tokens:
																	if skip_special_tokens and token in self.all_special_ids:
																						continue
																	if token in self.added_tokens_encoder:
																						if current_sub_text:
																											sub_texts.append(self.convert_tokens_to_string(__A    )    )
																											a__          =  []
																						sub_texts.append(__A    )
																	else:
																						current_sub_text.append(__A    )
												if current_sub_text:
																	sub_texts.append(self.convert_tokens_to_string(__A    )    )

												# Mimic the behavior of the Rust tokenizer:
												# No space before [MASK] and [SEP]
												if spaces_between_special_tokens:
																	a__          =  re.sub(R''' (\[(MASK|SEP)\])'''				,		R'''\1'''				,		''' '''.join(__A    )    )
												else:
																	a__          =  ''''''.join(__A    )

												a__          =  (
												    clean_up_tokenization_spaces
												    if clean_up_tokenization_spaces is not None
												    else self.clean_up_tokenization_spaces
												)
												if clean_up_tokenization_spaces:
																	a__          =  self.clean_up_tokenization(__A    )
																	return clean_text
												else:
																	return text




							def  lowercase		(   self:     Any				,		__A:     str				,		__A:     Optional[str] = None    ):

												'''simple docstring'''

												if not os.path.isdir(__A    ):
																	logger.error(F'Vocabulary path ({save_directory}) should be a directory'    )
																	return
												a__          =  os.path.join(
												    __A				,		(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']    )

												if os.path.abspath(self.vocab_file    ) != os.path.abspath(__A    ) and os.path.isfile(self.vocab_file    ):
																	copyfile(self.vocab_file				,		__A    )
												elif not os.path.isfile(self.vocab_file    ):
																	with open(__A				,		'''wb'''    ) as fi:
																						a__          =  self.sp_model.serialized_model_proto()
																						fi.write(__A    )

												return (out_vocab_file,)




							def  lowercase		(   self:     int				,		__A:     List[int]				,		__A:     Optional[List[int]] = None    ):

												'''simple docstring'''

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




							def  lowercase		(   self:     Union[str, Any]				,		__A:     List[int]				,		__A:     Optional[List[int]] = None				,		__A:     bool = False    ):

												'''simple docstring'''

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

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





							def  lowercase		(   self:     Tuple				,		__A:     List[int]				,		__A:     Optional[List[int]] = None    ):

												'''simple docstring'''

												a__          =  [self.sep_token_id]
												a__          =  [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]










"""simple docstring"""







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

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Union[str, Any]    ):

												'''simple docstring'''

												a__          =  []
												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=__A				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

																	a__          =  FlaxTransformeraDModel(
																	    in_channels=self.out_channels				,		n_heads=self.num_attention_heads				,		d_head=self.out_channels // self.num_attention_heads				,		depth=1				,		use_linear_projection=self.use_linear_projection				,		only_cross_attention=self.only_cross_attention				,		use_memory_efficient_attention=self.use_memory_efficient_attention				,		dtype=self.dtype				,		)
																	attentions.append(__A    )

												a__          =  resnets
												a__          =  attentions

												if self.add_downsample:
																	a__          =  FlaxDownsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Optional[int]				,		__A:     Union[str, Any]				,		__A:     str				,		__A:     Optional[Any]				,		__A:     Any=True    ):

												'''simple docstring'''

												a__          =  ()

												for resnet, attn in zip(self.resnets				,		self.attentions    ):
																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )
																	a__          =  attn(__A				,		__A				,		deterministic=__A    )
																	output_states += (hidden_states,)

												if self.add_downsample:
																	a__          =  self.downsamplers_a(__A    )
																	output_states += (hidden_states,)

												return hidden_states, output_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Union[str, Any]    ):

												'''simple docstring'''

												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=__A				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )
												a__          =  resnets

												if self.add_downsample:
																	a__          =  FlaxDownsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Dict				,		__A:     int				,		__A:     Dict				,		__A:     Optional[Any]=True    ):

												'''simple docstring'''

												a__          =  ()

												for resnet in self.resnets:
																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )
																	output_states += (hidden_states,)

												if self.add_downsample:
																	a__          =  self.downsamplers_a(__A    )
																	output_states += (hidden_states,)

												return hidden_states, output_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Optional[int]    ):

												'''simple docstring'''

												a__          =  []
												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if (i == self.num_layers - 1) else self.out_channels
																	a__          =  self.prev_output_channel if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=resnet_in_channels + res_skip_channels				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

																	a__          =  FlaxTransformeraDModel(
																	    in_channels=self.out_channels				,		n_heads=self.num_attention_heads				,		d_head=self.out_channels // self.num_attention_heads				,		depth=1				,		use_linear_projection=self.use_linear_projection				,		only_cross_attention=self.only_cross_attention				,		use_memory_efficient_attention=self.use_memory_efficient_attention				,		dtype=self.dtype				,		)
																	attentions.append(__A    )

												a__          =  resnets
												a__          =  attentions

												if self.add_upsample:
																	a__          =  FlaxUpsampleaD(self.out_channels				,		dtype=self.dtype    )





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

												'''simple docstring'''

												for resnet, attn in zip(self.resnets				,		self.attentions    ):
																	# pop res hidden states
																	a__          =  res_hidden_states_tuple[-1]
																	a__          =  res_hidden_states_tuple[:-1]
																	a__          =  jnp.concatenate((hidden_states, res_hidden_states)				,		axis=-1    )

																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )
																	a__          =  attn(__A				,		__A				,		deterministic=__A    )

												if self.add_upsample:
																	a__          =  self.upsamplers_a(__A    )

												return hidden_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     str    ):

												'''simple docstring'''

												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if (i == self.num_layers - 1) else self.out_channels
																	a__          =  self.prev_output_channel if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=resnet_in_channels + res_skip_channels				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

												a__          =  resnets

												if self.add_upsample:
																	a__          =  FlaxUpsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Tuple				,		__A:     Optional[Any]				,		__A:     Optional[Any]				,		__A:     Union[str, Any]				,		__A:     Dict=True    ):

												'''simple docstring'''

												for resnet in self.resnets:
																	# pop res hidden states
																	a__          =  res_hidden_states_tuple[-1]
																	a__          =  res_hidden_states_tuple[:-1]
																	a__          =  jnp.concatenate((hidden_states, res_hidden_states)				,		axis=-1    )

																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )

												if self.add_upsample:
																	a__          =  self.upsamplers_a(__A    )

												return hidden_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Tuple    ):

												'''simple docstring'''

												a__          =  [
												    FlaxResnetBlockaD(
												        in_channels=self.in_channels				,		out_channels=self.in_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
												]

												a__          =  []

												for _ in range(self.num_layers    ):
																	a__          =  FlaxTransformeraDModel(
																	    in_channels=self.in_channels				,		n_heads=self.num_attention_heads				,		d_head=self.in_channels // self.num_attention_heads				,		depth=1				,		use_linear_projection=self.use_linear_projection				,		use_memory_efficient_attention=self.use_memory_efficient_attention				,		dtype=self.dtype				,		)
																	attentions.append(__A    )

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=self.in_channels				,		out_channels=self.in_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

												a__          =  resnets
												a__          =  attentions





							def __call__(   self:     Any				,		__A:     Optional[int]				,		__A:     int				,		__A:     Tuple				,		__A:     str=True    ):

												'''simple docstring'''

												a__          =  self.resnets[0](__A				,		__A    )
												for attn, resnet in zip(self.attentions				,		self.resnets[1:]    ):
																	a__          =  attn(__A				,		__A				,		deterministic=__A    )
																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )

												return hidden_states


