
'''simple docstring'''






from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block

@dataclass
class 							_snake_case     (  __a ):
	"""simple docstring"""





	_UpperCamelCase							:  Dict								=		42







class 							_snake_case     (  __a						,    __a ):
	"""simple docstring"""





	@register_to_config
	def __init__(						self				: Dict					,						UpperCamelCase_				: Union[str, Any] = None					,						UpperCamelCase_				: List[str] = 3					,						UpperCamelCase_				: Optional[Any] = 3					,						UpperCamelCase_				: Optional[Any] = False					,						UpperCamelCase_				: str = "positional"					,						UpperCamelCase_				: List[Any] = 0					,						UpperCamelCase_				: str = True					,						UpperCamelCase_				: int = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")					,						UpperCamelCase_				: Optional[Any] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")					,						UpperCamelCase_				: Union[str, Any] = (224, 448, 672, 896)					,						UpperCamelCase_				: Tuple = 2					,						UpperCamelCase_				: int = 1					,						UpperCamelCase_				: Optional[Any] = 1					,						UpperCamelCase_				: Dict = "conv"					,						UpperCamelCase_				: Dict = "conv"					,						UpperCamelCase_				: Tuple = "silu"					,						UpperCamelCase_				: Dict = 8					,						UpperCamelCase_				: Tuple = 32					,						UpperCamelCase_				: Any = 1E-5					,						UpperCamelCase_				: str = "default"					,						UpperCamelCase_				: Any = True					,						UpperCamelCase_				: str = None					,						UpperCamelCase_				: str = None					,						):
			super().__init__()

			lowerCAmelCase_   :       List[str]    							=sample_size
			lowerCAmelCase_   :       Tuple    							=block_out_channels[0] * 4

			# Check inputs
			if len(SCREAMING_SNAKE_CASE__      ) != len(SCREAMING_SNAKE_CASE__      ):
					raise ValueError(
					    F'Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}.'      )

			if len(SCREAMING_SNAKE_CASE__      ) != len(SCREAMING_SNAKE_CASE__      ):
					raise ValueError(
					    F'Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}.'      )

			# input
			lowerCAmelCase_   :       Any    							=nn.Convad(SCREAMING_SNAKE_CASE__					,						block_out_channels[0]					,						kernel_size=3					,						padding=(1, 1)      )

			# time
			if time_embedding_type == "fourier":
					lowerCAmelCase_   :       Union[str, Any]    							=GaussianFourierProjection(embedding_size=block_out_channels[0]					,						scale=16      )
					lowerCAmelCase_   :       str    							=2 * block_out_channels[0]
			elif time_embedding_type == "positional":
					lowerCAmelCase_   :       List[str]    							=Timesteps(block_out_channels[0]					,						SCREAMING_SNAKE_CASE__					,						SCREAMING_SNAKE_CASE__      )
					lowerCAmelCase_   :       Tuple    							=block_out_channels[0]

			lowerCAmelCase_   :       Optional[Any]    							=TimestepEmbedding(SCREAMING_SNAKE_CASE__					,						SCREAMING_SNAKE_CASE__      )

			# class embedding
			if class_embed_type is None and num_class_embeds is not None:
					lowerCAmelCase_   :       Dict    							=nn.Embedding(SCREAMING_SNAKE_CASE__					,						SCREAMING_SNAKE_CASE__      )
			elif class_embed_type == "timestep":
					lowerCAmelCase_   :       str    							=TimestepEmbedding(SCREAMING_SNAKE_CASE__					,						SCREAMING_SNAKE_CASE__      )
			elif class_embed_type == "identity":
					lowerCAmelCase_   :       List[str]    							=nn.Identity(SCREAMING_SNAKE_CASE__					,						SCREAMING_SNAKE_CASE__      )
			else:
					lowerCAmelCase_   :       Dict    							=None

			lowerCAmelCase_   :       Optional[Any]    							=nn.ModuleList([]      )
			lowerCAmelCase_   :       List[str]    							=None
			lowerCAmelCase_   :       Optional[int]    							=nn.ModuleList([]      )

			# down
			lowerCAmelCase_   :       int    							=block_out_channels[0]
			for i, down_block_type in enumerate(SCREAMING_SNAKE_CASE__      ):
					lowerCAmelCase_   :       Dict    							=output_channel
					lowerCAmelCase_   :       Union[str, Any]    							=block_out_channels[i]
					lowerCAmelCase_   :       Any    							=i == len(SCREAMING_SNAKE_CASE__      ) - 1

					lowerCAmelCase_   :       Dict    							=get_down_block(
					    SCREAMING_SNAKE_CASE__					,						num_layers=SCREAMING_SNAKE_CASE__					,						in_channels=SCREAMING_SNAKE_CASE__					,						out_channels=SCREAMING_SNAKE_CASE__					,						temb_channels=SCREAMING_SNAKE_CASE__					,						add_downsample=not is_final_block					,						resnet_eps=SCREAMING_SNAKE_CASE__					,						resnet_act_fn=SCREAMING_SNAKE_CASE__					,						resnet_groups=SCREAMING_SNAKE_CASE__					,						attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel					,						downsample_padding=SCREAMING_SNAKE_CASE__					,						resnet_time_scale_shift=SCREAMING_SNAKE_CASE__					,						downsample_type=SCREAMING_SNAKE_CASE__					,						)
					self.down_blocks.append(SCREAMING_SNAKE_CASE__      )

			# mid
			lowerCAmelCase_   :       Union[str, Any]    							=UNetMidBlockaD(
			    in_channels=block_out_channels[-1]					,						temb_channels=SCREAMING_SNAKE_CASE__					,						resnet_eps=SCREAMING_SNAKE_CASE__					,						resnet_act_fn=SCREAMING_SNAKE_CASE__					,						output_scale_factor=SCREAMING_SNAKE_CASE__					,						resnet_time_scale_shift=SCREAMING_SNAKE_CASE__					,						attention_head_dim=attention_head_dim if attention_head_dim is not None else block_out_channels[-1]					,						resnet_groups=SCREAMING_SNAKE_CASE__					,						add_attention=SCREAMING_SNAKE_CASE__					,						)

			# up
			lowerCAmelCase_   :       List[Any]    							=list(reversed(SCREAMING_SNAKE_CASE__      )      )
			lowerCAmelCase_   :       Tuple    							=reversed_block_out_channels[0]
			for i, up_block_type in enumerate(SCREAMING_SNAKE_CASE__      ):
					lowerCAmelCase_   :       Union[str, Any]    							=output_channel
					lowerCAmelCase_   :       Union[str, Any]    							=reversed_block_out_channels[i]
					lowerCAmelCase_   :       int    							=reversed_block_out_channels[min(i + 1					,						len(SCREAMING_SNAKE_CASE__      ) - 1      )]

					lowerCAmelCase_   :       int    							=i == len(SCREAMING_SNAKE_CASE__      ) - 1

					lowerCAmelCase_   :       Union[str, Any]    							=get_up_block(
					    SCREAMING_SNAKE_CASE__					,						num_layers=layers_per_block + 1					,						in_channels=SCREAMING_SNAKE_CASE__					,						out_channels=SCREAMING_SNAKE_CASE__					,						prev_output_channel=SCREAMING_SNAKE_CASE__					,						temb_channels=SCREAMING_SNAKE_CASE__					,						add_upsample=not is_final_block					,						resnet_eps=SCREAMING_SNAKE_CASE__					,						resnet_act_fn=SCREAMING_SNAKE_CASE__					,						resnet_groups=SCREAMING_SNAKE_CASE__					,						attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel					,						resnet_time_scale_shift=SCREAMING_SNAKE_CASE__					,						upsample_type=SCREAMING_SNAKE_CASE__					,						)
					self.up_blocks.append(SCREAMING_SNAKE_CASE__      )
					lowerCAmelCase_   :       int    							=output_channel

			# out
			lowerCAmelCase_   :       str    							=norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4					,						32      )
			lowerCAmelCase_   :       str    							=nn.GroupNorm(num_channels=block_out_channels[0]					,						num_groups=SCREAMING_SNAKE_CASE__					,						eps=SCREAMING_SNAKE_CASE__      )
			lowerCAmelCase_   :       Optional[Any]    							=nn.SiLU()
			lowerCAmelCase_   :       List[str]    							=nn.Convad(block_out_channels[0]					,						SCREAMING_SNAKE_CASE__					,						kernel_size=3					,						padding=1      )





	def 	__A      (						self				: Tuple					,						UpperCamelCase_				: int					,						UpperCamelCase_				: Union[str, Any]					,						UpperCamelCase_				: Dict = None					,						UpperCamelCase_				: Dict = True					,						):
			# 0. center input if necessary
			if self.config.center_input_sample:
					lowerCAmelCase_   :       Any    							=2 * sample - 1.0

			# 1. time
			lowerCAmelCase_   :       Any    							=timestep
			if not torch.is_tensor(SCREAMING_SNAKE_CASE__      ):
					lowerCAmelCase_   :       Tuple    							=torch.tensor([timesteps]					,						dtype=torch.long					,						device=sample.device      )
			elif torch.is_tensor(SCREAMING_SNAKE_CASE__      ) and len(timesteps.shape      ) == 0:
					lowerCAmelCase_   :       Union[str, Any]    							=timesteps[None].to(sample.device      )

			# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
			lowerCAmelCase_   :       Optional[int]    							=timesteps * torch.ones(sample.shape[0]					,						dtype=timesteps.dtype					,						device=timesteps.device      )

			lowerCAmelCase_   :       str    							=self.time_proj(SCREAMING_SNAKE_CASE__      )

			# timesteps does not contain any weights and will always return f32 tensors
			# but time_embedding might actually be running in fp16. so we need to cast here.
			# there might be better ways to encapsulate this.
			lowerCAmelCase_   :       Optional[int]    							=t_emb.to(dtype=self.dtype      )
			lowerCAmelCase_   :       Optional[int]    							=self.time_embedding(SCREAMING_SNAKE_CASE__      )

			if self.class_embedding is not None:
					if class_labels is None:
							raise ValueError('''class_labels should be provided when doing class conditioning'''      )

					if self.config.class_embed_type == "timestep":
							lowerCAmelCase_   :       Union[str, Any]    							=self.time_proj(SCREAMING_SNAKE_CASE__      )

					lowerCAmelCase_   :       Optional[int]    							=self.class_embedding(SCREAMING_SNAKE_CASE__      ).to(dtype=self.dtype      )
					lowerCAmelCase_   :       Tuple    							=emb + class_emb

			# 2. pre-process
			lowerCAmelCase_   :       List[str]    							=sample
			lowerCAmelCase_   :       Dict    							=self.conv_in(SCREAMING_SNAKE_CASE__      )

			# 3. down
			lowerCAmelCase_   :       Optional[Any]    							=(sample,)
			for downsample_block in self.down_blocks:
					if hasattr(SCREAMING_SNAKE_CASE__					,						'''skip_conv'''      ):
							lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_   :       Optional[Any]    							=downsample_block(
							    hidden_states=SCREAMING_SNAKE_CASE__					,						temb=SCREAMING_SNAKE_CASE__					,						skip_sample=SCREAMING_SNAKE_CASE__      )
					else:
							lowerCAmelCase_							,		lowerCAmelCase_   :       List[str]    							=downsample_block(hidden_states=SCREAMING_SNAKE_CASE__					,						temb=SCREAMING_SNAKE_CASE__      )

					down_block_res_samples += res_samples

			# 4. mid
			lowerCAmelCase_   :       Optional[int]    							=self.mid_block(SCREAMING_SNAKE_CASE__					,						SCREAMING_SNAKE_CASE__      )

			# 5. up
			lowerCAmelCase_   :       List[Any]    							=None
			for upsample_block in self.up_blocks:
					lowerCAmelCase_   :       List[str]    							=down_block_res_samples[-len(upsample_block.resnets      ) :]
					lowerCAmelCase_   :       int    							=down_block_res_samples[: -len(upsample_block.resnets      )]

					if hasattr(SCREAMING_SNAKE_CASE__					,						'''skip_conv'''      ):
							lowerCAmelCase_							,		lowerCAmelCase_   :       Dict    							=upsample_block(SCREAMING_SNAKE_CASE__					,						SCREAMING_SNAKE_CASE__					,						SCREAMING_SNAKE_CASE__					,						SCREAMING_SNAKE_CASE__      )
					else:
							lowerCAmelCase_   :       Any    							=upsample_block(SCREAMING_SNAKE_CASE__					,						SCREAMING_SNAKE_CASE__					,						SCREAMING_SNAKE_CASE__      )

        # 6. post-process
			lowerCAmelCase_   :       Dict    							=self.conv_norm_out(SCREAMING_SNAKE_CASE__      )
			lowerCAmelCase_   :       int    							=self.conv_act(SCREAMING_SNAKE_CASE__      )
			lowerCAmelCase_   :       Optional[Any]    							=self.conv_out(SCREAMING_SNAKE_CASE__      )

			if skip_sample is not None:
					sample += skip_sample

			if self.config.time_embedding_type == "fourier":
					lowerCAmelCase_   :       Union[str, Any]    							=timesteps.reshape((sample.shape[0], *([1] * len(sample.shape[1:]      )))      )
					lowerCAmelCase_   :       int    							=sample / timesteps

			if not return_dict:
					return (sample,)

			return UNetaDOutput(sample=SCREAMING_SNAKE_CASE__      )









"""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_splinter import SplinterTokenizer


SCREAMING_SNAKE_CASE_  :      str							=      logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_  :      List[Any]							=      {'''vocab_file''': '''vocab.txt'''}

SCREAMING_SNAKE_CASE_  :      int							=      {
    '''vocab_file''': {
        '''tau/splinter-base''': '''https://huggingface.co/tau/splinter-base/resolve/main/vocab.txt''',
        '''tau/splinter-base-qass''': '''https://huggingface.co/tau/splinter-base-qass/resolve/main/vocab.txt''',
        '''tau/splinter-large''': '''https://huggingface.co/tau/splinter-large/resolve/main/vocab.txt''',
        '''tau/splinter-large-qass''': '''https://huggingface.co/tau/splinter-large-qass/resolve/main/vocab.txt''',
    }
}

SCREAMING_SNAKE_CASE_  :      Union[str, Any]							=      {
    '''tau/splinter-base''': 512,
    '''tau/splinter-base-qass''': 512,
    '''tau/splinter-large''': 512,
    '''tau/splinter-large-qass''': 512,
}

SCREAMING_SNAKE_CASE_  :      List[Any]							=      {
    '''tau/splinter-base''': {'''do_lower_case''': False},
    '''tau/splinter-base-qass''': {'''do_lower_case''': False},
    '''tau/splinter-large''': {'''do_lower_case''': False},
    '''tau/splinter-large-qass''': {'''do_lower_case''': False},
}


class 		_A							(      __a       ):
      __a             =						VOCAB_FILES_NAMES
      __a             =						PRETRAINED_VOCAB_FILES_MAP
      __a             =						PRETRAINED_INIT_CONFIGURATION
      __a             =						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __a             =						SplinterTokenizer





      def __init__(	self     ,    SCREAMING_SNAKE_CASE__=None     ,    SCREAMING_SNAKE_CASE__=None     ,    SCREAMING_SNAKE_CASE__=True     ,    SCREAMING_SNAKE_CASE__="[UNK]"     ,    SCREAMING_SNAKE_CASE__="[SEP]"     ,    SCREAMING_SNAKE_CASE__="[PAD]"     ,    SCREAMING_SNAKE_CASE__="[CLS]"     ,    SCREAMING_SNAKE_CASE__="[MASK]"     ,    SCREAMING_SNAKE_CASE__="[QUESTION]"     ,    SCREAMING_SNAKE_CASE__=True     ,    SCREAMING_SNAKE_CASE__=None     ,    **SCREAMING_SNAKE_CASE__     ,    ) ->       Optional[Any]:
         super().__init__(
             SCREAMING_SNAKE_CASE__     ,    tokenizer_file=SCREAMING_SNAKE_CASE__     ,    do_lower_case=SCREAMING_SNAKE_CASE__     ,    unk_token=SCREAMING_SNAKE_CASE__     ,    sep_token=SCREAMING_SNAKE_CASE__     ,    pad_token=SCREAMING_SNAKE_CASE__     ,    cls_token=SCREAMING_SNAKE_CASE__     ,    mask_token=SCREAMING_SNAKE_CASE__     ,    tokenize_chinese_chars=SCREAMING_SNAKE_CASE__     ,    strip_accents=SCREAMING_SNAKE_CASE__     ,    additional_special_tokens=(question_token,)     ,    **SCREAMING_SNAKE_CASE__     ,    )

         lowerCamelCase__						       =      json.loads(self.backend_tokenizer.normalizer.__getstate__()		)
         if (
             pre_tok_state.get("lowercase"     ,    SCREAMING_SNAKE_CASE__		) != do_lower_case
             or pre_tok_state.get("strip_accents"     ,    SCREAMING_SNAKE_CASE__		) != strip_accents
         ):
            lowerCamelCase__						       =      getattr(SCREAMING_SNAKE_CASE__     ,    pre_tok_state.pop("type"		)		)
            lowerCamelCase__						       =      do_lower_case
            lowerCamelCase__						       =      strip_accents
            lowerCamelCase__						       =      pre_tok_class(**SCREAMING_SNAKE_CASE__		)

         lowerCamelCase__						       =      do_lower_case





      @property
      def        _lowerCamelCase	(	self		) ->       Optional[Any]:
         return self.convert_tokens_to_ids(self.question_token		)





      def        _lowerCamelCase	(	self     ,    SCREAMING_SNAKE_CASE__     ,    SCREAMING_SNAKE_CASE__ = None		) ->       List[int]:
         if token_ids_a is None:
            return [self.cls_token_id] + token_ids_a + [self.sep_token_id]

         lowerCamelCase__						       =      [self.cls_token_id]
         lowerCamelCase__						       =      [self.sep_token_id]
         lowerCamelCase__						       =      [self.question_token_id] + [self.convert_tokens_to_ids("."		)]
         if self.padding_side == "right":
            # Input is question-then-context
            return cls + token_ids_a + question_suffix + sep + token_ids_a + sep
         else:
            # Input is context-then-question
            return cls + token_ids_a + sep + token_ids_a + question_suffix + sep





      def        _lowerCamelCase	(	self     ,    SCREAMING_SNAKE_CASE__     ,    SCREAMING_SNAKE_CASE__ = None		) ->       List[int]:
         lowerCamelCase__						       =      [self.sep_token_id]
         lowerCamelCase__						       =      [self.cls_token_id]
         lowerCamelCase__						       =      [self.question_token_id] + [self.convert_tokens_to_ids("."		)]
         if token_ids_a is None:
            return len(cls + token_ids_a + sep		) * [0]

         if self.padding_side == "right":
            # Input is question-then-context
            return len(cls + token_ids_a + question_suffix + sep		) * [0] + len(token_ids_a + sep		) * [1]
         else:
            # Input is context-then-question
            return len(cls + token_ids_a + sep		) * [0] + len(token_ids_a + question_suffix + sep		) * [1]





      def        _lowerCamelCase	(	self     ,    SCREAMING_SNAKE_CASE__     ,    SCREAMING_SNAKE_CASE__ = None		) ->       Tuple[str]:
         lowerCamelCase__						       =      self._tokenizer.model.save(SCREAMING_SNAKE_CASE__     ,    name=SCREAMING_SNAKE_CASE__		)
         return tuple(SCREAMING_SNAKE_CASE__		)




