






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


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

_UpperCamelCase			:	Any								={'vocab_file': 'spiece.model'}

_UpperCamelCase			:	Dict								={
    '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'
        ),
    }
}

_UpperCamelCase			:	Any								={
    'google/bigbird-roberta-base': 4096,
    'google/bigbird-roberta-large': 4096,
    'google/bigbird-base-trivia-itc': 4096,
}






class 						UpperCAmelCase__ (							__snake_case							):
		__snake_case      :       str     =   VOCAB_FILES_NAMES
		__snake_case      :       Optional[int]     =   PRETRAINED_VOCAB_FILES_MAP
		__snake_case      :       Optional[Any]     =   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
		__snake_case      :       List[Any]     =   ["input_ids", "attention_mask"]
		__snake_case      :       List[int]     =   []







		def __init__(					self       ,A__       ,A__="<unk>"       ,A__="<s>"       ,A__="</s>"       ,A__="<pad>"       ,A__="[SEP]"       ,A__="[MASK]"       ,A__="[CLS]"       ,A__ = None       ,**A__       ,):
						_A	:				Dict   =      AddedToken(A__       ,lstrip=A__       ,rstrip=A__		) if isinstance(A__       ,A__		) else bos_token
						_A	:				Any   =      AddedToken(A__       ,lstrip=A__       ,rstrip=A__		) if isinstance(A__       ,A__		) else eos_token
						_A	:				List[str]   =      AddedToken(A__       ,lstrip=A__       ,rstrip=A__		) if isinstance(A__       ,A__		) else unk_token
						_A	:				int   =      AddedToken(A__       ,lstrip=A__       ,rstrip=A__		) if isinstance(A__       ,A__		) else pad_token
						_A	:				str   =      AddedToken(A__       ,lstrip=A__       ,rstrip=A__		) if isinstance(A__       ,A__		) else cls_token
						_A	:				int   =      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	:				Optional[int]   =      AddedToken(A__       ,lstrip=A__       ,rstrip=A__		) if isinstance(A__       ,A__		) else mask_token

						_A	:				Union[str, Any]   =      {} 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	:				Any   =      vocab_file

						_A	:				List[str]   =      spm.SentencePieceProcessor(**self.sp_model_kwargs		)
						self.sp_model.Load(A__		)







		@property
		def  A__      (					self		):
						return self.sp_model.get_piece_size()







		def  A__      (					self		):
						_A	:				Optional[int]   =      {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		):
						_A	:				List[str]   =      self.__dict__.copy()
						_A	:				Optional[int]   =      None
						return state







		def __setstate__(					self       ,A__		):
						_A	:				List[Any]   =      d

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

						_A	:				List[str]   =      spm.SentencePieceProcessor(**self.sp_model_kwargs		)
						self.sp_model.Load(self.vocab_file		)







		def  A__      (					self       ,A__		):
						return self.sp_model.encode(A__       ,out_type=A__		)







		def  A__      (					self       ,A__		):
						return self.sp_model.piece_to_id(A__		)







		def  A__      (					self       ,A__		):
						_A	:				int   =      self.sp_model.IdToPiece(A__		)
						return token







		def  A__      (					self       ,A__		):
						_A	:				List[Any]   =      []
						_A	:				int   =      ''''''
						_A	:				Any   =      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	:				List[Any]   =      True
														_A	:				List[str]   =      []
										else:
														current_sub_tokens.append(A__		)
														_A	:				Optional[int]   =      False
						out_string += self.sp_model.decode(A__		)
						return out_string.strip()







		def  A__      (					self       ,A__       ,A__ = False       ,A__ = None       ,A__ = True       ,**A__       ,):
						_A	:				List[Any]   =      kwargs.pop('''use_source_tokenizer'''       ,A__		)

						_A	:				str   =      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	:				List[str]   =      []
						_A	:				int   =      []
						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	:				Optional[int]   =      []
														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	:				List[Any]   =      re.sub(r''' (\[(MASK|SEP)\])'''       ,r'''\1'''       ,''' '''.join(A__		)		)
						else:
										_A	:				List[str]   =      ''''''.join(A__		)

						_A	:				Any   =      (
						    clean_up_tokenization_spaces
						    if clean_up_tokenization_spaces is not None
						    else self.clean_up_tokenization_spaces
						)
						if clean_up_tokenization_spaces:
										_A	:				List[str]   =      self.clean_up_tokenization(A__		)
										return clean_text
						else:
										return text







		def  A__      (					self       ,A__       ,A__ = None		):
						if not os.path.isdir(A__		):
										logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""		)
										return
						_A	:				Tuple   =      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	:				Optional[Any]   =      self.sp_model.serialized_model_proto()
														fi.write(A__		)

						return (out_vocab_file,)







		def  A__      (					self       ,A__       ,A__ = None		):
						if token_ids_a is None:
										return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
						_A	:				str   =      [self.cls_token_id]
						_A	:				List[Any]   =      [self.sep_token_id]
						return cls + token_ids_a + sep + token_ids_a + sep







		def  A__      (					self       ,A__       ,A__ = None       ,A__ = False		):
						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  A__      (					self       ,A__       ,A__ = None		):
						_A	:				Union[str, Any]   =      [self.sep_token_id]
						_A	:				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]











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

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






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				Optional[Any]   =      []
						_A	:				str   =      []

						for i in range(self.num_layers		):
										_A	:				Union[str, Any]   =      self.in_channels if i == 0 else self.out_channels

										_A	:				str   =      FlaxResnetBlockaD(
										    in_channels=A__       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

										_A	:				List[str]   =      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	:				List[str]   =      resnets
						_A	:				Any   =      attentions

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







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						_A	:				List[str]   =      ()

						for resnet, attn in zip(self.resnets       ,self.attentions		):
										_A	:				Optional[Any]   =      resnet(A__       ,A__       ,deterministic=A__		)
										_A	:				Tuple   =      attn(A__       ,A__       ,deterministic=A__		)
										output_states += (hidden_states,)

						if self.add_downsample:
										_A	:				Optional[int]   =      self.downsamplers_a(A__		)
										output_states += (hidden_states,)

						return hidden_states, output_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				List[Any]   =      []

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

										_A	:				List[Any]   =      FlaxResnetBlockaD(
										    in_channels=A__       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)
						_A	:				Tuple   =      resnets

						if self.add_downsample:
										_A	:				Optional[Any]   =      FlaxDownsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__=True		):
						_A	:				List[Any]   =      ()

						for resnet in self.resnets:
										_A	:				Optional[int]   =      resnet(A__       ,A__       ,deterministic=A__		)
										output_states += (hidden_states,)

						if self.add_downsample:
										_A	:				List[Any]   =      self.downsamplers_a(A__		)
										output_states += (hidden_states,)

						return hidden_states, output_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				List[str]   =      []
						_A	:				Optional[Any]   =      []

						for i in range(self.num_layers		):
										_A	:				Dict   =      self.in_channels if (i == self.num_layers - 1) else self.out_channels
										_A	:				List[str]   =      self.prev_output_channel if i == 0 else self.out_channels

										_A	:				Optional[int]   =      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	:				Union[str, Any]   =      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	:				Dict   =      resnets
						_A	:				int   =      attentions

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







		def __call__(					self       ,A__       ,A__       ,A__       ,A__       ,A__=True		):
						for resnet, attn in zip(self.resnets       ,self.attentions		):
										# pop res hidden states
										_A	:				List[Any]   =      res_hidden_states_tuple[-1]
										_A	:				int   =      res_hidden_states_tuple[:-1]
										_A	:				List[Any]   =      jnp.concatenate((hidden_states, res_hidden_states)       ,axis=-1		)

										_A	:				str   =      resnet(A__       ,A__       ,deterministic=A__		)
										_A	:				Optional[int]   =      attn(A__       ,A__       ,deterministic=A__		)

						if self.add_upsample:
										_A	:				Union[str, Any]   =      self.upsamplers_a(A__		)

						return hidden_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				Optional[Any]   =      []

						for i in range(self.num_layers		):
										_A	:				Any   =      self.in_channels if (i == self.num_layers - 1) else self.out_channels
										_A	:				Optional[Any]   =      self.prev_output_channel if i == 0 else self.out_channels

										_A	:				List[Any]   =      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	:				int   =      resnets

						if self.add_upsample:
										_A	:				List[str]   =      FlaxUpsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						for resnet in self.resnets:
										# pop res hidden states
										_A	:				Tuple   =      res_hidden_states_tuple[-1]
										_A	:				Optional[int]   =      res_hidden_states_tuple[:-1]
										_A	:				Any   =      jnp.concatenate((hidden_states, res_hidden_states)       ,axis=-1		)

										_A	:				Optional[Any]   =      resnet(A__       ,A__       ,deterministic=A__		)

						if self.add_upsample:
										_A	:				Tuple   =      self.upsamplers_a(A__		)

						return hidden_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						# there is always at least one resnet
						_A	:				List[str]   =      [
						    FlaxResnetBlockaD(
						        in_channels=self.in_channels       ,out_channels=self.in_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
						]

						_A	:				Union[str, Any]   =      []

						for _ in range(self.num_layers		):
										_A	:				Any   =      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	:				Any   =      FlaxResnetBlockaD(
										    in_channels=self.in_channels       ,out_channels=self.in_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

						_A	:				List[str]   =      resnets
						_A	:				str   =      attentions







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						_A	:				Optional[int]   =      self.resnets[0](A__       ,A__		)
						for attn, resnet in zip(self.attentions       ,self.resnets[1:]		):
										_A	:				Dict   =      attn(A__       ,A__       ,deterministic=A__		)
										_A	:				Any   =      resnet(A__       ,A__       ,deterministic=A__		)

						return hidden_states



