



'''simple docstring'''



import json
from typing import List, Optional, Tuple

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_roberta import RobertaTokenizer


lowerCAmelCase	:		str				=   logging.get_logger(__name__)

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

lowerCAmelCase	:		Optional[int]				=   {
    '''vocab_file''': {
        '''roberta-base''': '''https://huggingface.co/roberta-base/resolve/main/vocab.json''',
        '''roberta-large''': '''https://huggingface.co/roberta-large/resolve/main/vocab.json''',
        '''roberta-large-mnli''': '''https://huggingface.co/roberta-large-mnli/resolve/main/vocab.json''',
        '''distilroberta-base''': '''https://huggingface.co/distilroberta-base/resolve/main/vocab.json''',
        '''roberta-base-openai-detector''': '''https://huggingface.co/roberta-base-openai-detector/resolve/main/vocab.json''',
        '''roberta-large-openai-detector''': (
            '''https://huggingface.co/roberta-large-openai-detector/resolve/main/vocab.json'''
        ),
    },
    '''merges_file''': {
        '''roberta-base''': '''https://huggingface.co/roberta-base/resolve/main/merges.txt''',
        '''roberta-large''': '''https://huggingface.co/roberta-large/resolve/main/merges.txt''',
        '''roberta-large-mnli''': '''https://huggingface.co/roberta-large-mnli/resolve/main/merges.txt''',
        '''distilroberta-base''': '''https://huggingface.co/distilroberta-base/resolve/main/merges.txt''',
        '''roberta-base-openai-detector''': '''https://huggingface.co/roberta-base-openai-detector/resolve/main/merges.txt''',
        '''roberta-large-openai-detector''': (
            '''https://huggingface.co/roberta-large-openai-detector/resolve/main/merges.txt'''
        ),
    },
    '''tokenizer_file''': {
        '''roberta-base''': '''https://huggingface.co/roberta-base/resolve/main/tokenizer.json''',
        '''roberta-large''': '''https://huggingface.co/roberta-large/resolve/main/tokenizer.json''',
        '''roberta-large-mnli''': '''https://huggingface.co/roberta-large-mnli/resolve/main/tokenizer.json''',
        '''distilroberta-base''': '''https://huggingface.co/distilroberta-base/resolve/main/tokenizer.json''',
        '''roberta-base-openai-detector''': (
            '''https://huggingface.co/roberta-base-openai-detector/resolve/main/tokenizer.json'''
        ),
        '''roberta-large-openai-detector''': (
            '''https://huggingface.co/roberta-large-openai-detector/resolve/main/tokenizer.json'''
        ),
    },
}

lowerCAmelCase	:		str				=   {
    '''roberta-base''': 5_1_2,
    '''roberta-large''': 5_1_2,
    '''roberta-large-mnli''': 5_1_2,
    '''distilroberta-base''': 5_1_2,
    '''roberta-base-openai-detector''': 5_1_2,
    '''roberta-large-openai-detector''': 5_1_2,
}






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



					def __init__(    self			,				UpperCamelCase=None			,				UpperCamelCase=None			,				UpperCamelCase=None			,				UpperCamelCase="replace"			,				UpperCamelCase="<s>"			,				UpperCamelCase="</s>"			,				UpperCamelCase="</s>"			,				UpperCamelCase="<s>"			,				UpperCamelCase="<unk>"			,				UpperCamelCase="<pad>"			,				UpperCamelCase="<mask>"			,				UpperCamelCase=False			,				UpperCamelCase=True			,				**UpperCamelCase			,				)     ->			Tuple:
						super().__init__(
						    UpperCamelCase			,				UpperCamelCase			,				tokenizer_file=UpperCamelCase			,				errors=UpperCamelCase			,				bos_token=UpperCamelCase			,				eos_token=UpperCamelCase			,				sep_token=UpperCamelCase			,				cls_token=UpperCamelCase			,				unk_token=UpperCamelCase			,				pad_token=UpperCamelCase			,				mask_token=UpperCamelCase			,				add_prefix_space=UpperCamelCase			,				trim_offsets=UpperCamelCase			,				**UpperCamelCase			,				)

						__lowerCAmelCase												=  json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()     )
						if pre_tok_state.get("add_prefix_space"			,				UpperCamelCase     ) != add_prefix_space:
							__lowerCAmelCase												=  getattr(UpperCamelCase			,				pre_tok_state.pop("type"     )     )
							__lowerCAmelCase												=  add_prefix_space
							__lowerCAmelCase												=  pre_tok_class(**UpperCamelCase     )

						__lowerCAmelCase												=  add_prefix_space

						__lowerCAmelCase												=  "post_processor"
						__lowerCAmelCase												=  getattr(self.backend_tokenizer			,				UpperCamelCase			,				UpperCamelCase     )
						if tokenizer_component_instance:
							__lowerCAmelCase												=  json.loads(tokenizer_component_instance.__getstate__()     )

							# The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
							if "sep" in state:
								__lowerCAmelCase												=  tuple(state["sep"]     )
							if "cls" in state:
								__lowerCAmelCase												=  tuple(state["cls"]     )

							__lowerCAmelCase												=  False

							if state.get("add_prefix_space"			,				UpperCamelCase     ) != add_prefix_space:
								__lowerCAmelCase												=  add_prefix_space
								__lowerCAmelCase												=  True

							if state.get("trim_offsets"			,				UpperCamelCase     ) != trim_offsets:
								__lowerCAmelCase												=  trim_offsets
								__lowerCAmelCase												=  True

							if changes_to_apply:
								__lowerCAmelCase												=  getattr(UpperCamelCase			,				state.pop("type"     )     )
								__lowerCAmelCase												=  component_class(**UpperCamelCase     )
								setattr(self.backend_tokenizer			,				UpperCamelCase			,				UpperCamelCase     )



					@property
					def        UpperCAmelCase_ (    self     )     ->			str:
						if self._mask_token is None:
							if self.verbose:
								logger.error("Using mask_token, but it is not set yet."     )
							return None
						return str(self._mask_token     )



					@mask_token.setter
					def        UpperCAmelCase_ (    self			,				UpperCamelCase     )     ->			int:
						__lowerCAmelCase												=  AddedToken(UpperCamelCase			,				lstrip=UpperCamelCase			,				rstrip=UpperCamelCase     ) if isinstance(UpperCamelCase			,				UpperCamelCase     ) else value
						__lowerCAmelCase												=  value



					def        UpperCAmelCase_ (    self			,				*UpperCamelCase			,				**UpperCamelCase     )     ->			BatchEncoding:
						__lowerCAmelCase												=  kwargs.get("is_split_into_words"			,				UpperCamelCase     )
						assert self.add_prefix_space or not is_split_into_words, (
						    F'''You need to instantiate {self.__class__.__name__} with add_prefix_space=True '''
						    "to use it with pretokenized inputs."
						)

						return super()._batch_encode_plus(*UpperCamelCase			,				**UpperCamelCase     )



					def        UpperCAmelCase_ (    self			,				*UpperCamelCase			,				**UpperCamelCase     )     ->			BatchEncoding:
						__lowerCAmelCase												=  kwargs.get("is_split_into_words"			,				UpperCamelCase     )

						assert self.add_prefix_space or not is_split_into_words, (
						    F'''You need to instantiate {self.__class__.__name__} with add_prefix_space=True '''
						    "to use it with pretokenized inputs."
						)

						return super()._encode_plus(*UpperCamelCase			,				**UpperCamelCase     )



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase = None     )     ->			Tuple[str]:
						__lowerCAmelCase												=  self._tokenizer.model.save(UpperCamelCase			,				name=UpperCamelCase     )
						return tuple(UpperCamelCase     )



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase=None     )     ->			Dict:
						__lowerCAmelCase												=  [self.bos_token_id] + token_ids_a + [self.eos_token_id]
						if token_ids_a is None:
							return output

						return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase = None     )     ->			List[int]:
						__lowerCAmelCase												=  [self.sep_token_id]
						__lowerCAmelCase												=  [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 + sep + token_ids_a + sep     ) * [0]




'''simple docstring'''



from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin






@dataclass
class 					UpperCAmelCase__	(		UpperCamelCase__ ):
					a      :     torch.FloatTensor






class 					UpperCAmelCase__	(		UpperCamelCase__							,       UpperCamelCase__ ):



					@register_to_config
					def __init__(    self			,				UpperCamelCase = 16			,				UpperCamelCase = 88			,				UpperCamelCase = None			,				UpperCamelCase = None			,				UpperCamelCase = 1			,				UpperCamelCase = 0.0			,				UpperCamelCase = 32			,				UpperCamelCase = None			,				UpperCamelCase = False			,				UpperCamelCase = None			,				UpperCamelCase = "geglu"			,				UpperCamelCase = True			,				UpperCamelCase = True			,				)     ->			List[str]:
						super().__init__()
						__lowerCAmelCase												=  num_attention_heads
						__lowerCAmelCase												=  attention_head_dim
						__lowerCAmelCase												=  num_attention_heads * attention_head_dim

						__lowerCAmelCase												=  in_channels

						__lowerCAmelCase												=  torch.nn.GroupNorm(num_groups=UpperCamelCase			,				num_channels=UpperCamelCase			,				eps=1E-6			,				affine=UpperCamelCase     )
						__lowerCAmelCase												=  nn.Linear(UpperCamelCase			,				UpperCamelCase     )

						# 3. Define transformers blocks
						__lowerCAmelCase												=  nn.ModuleList(
						    [
						        BasicTransformerBlock(
						            UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				dropout=UpperCamelCase			,				cross_attention_dim=UpperCamelCase			,				activation_fn=UpperCamelCase			,				attention_bias=UpperCamelCase			,				double_self_attention=UpperCamelCase			,				norm_elementwise_affine=UpperCamelCase			,				)
						        for d in range(UpperCamelCase     )
						    ]     )

						__lowerCAmelCase												=  nn.Linear(UpperCamelCase			,				UpperCamelCase     )



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase=None			,				UpperCamelCase=None			,				UpperCamelCase=None			,				UpperCamelCase=1			,				UpperCamelCase=None			,				UpperCamelCase = True			,				)     ->			List[str]:
						__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase												=  hidden_states.shape
						__lowerCAmelCase												=  batch_frames // num_frames

						__lowerCAmelCase												=  hidden_states

						__lowerCAmelCase												=  hidden_states[None, :].reshape(UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )
						__lowerCAmelCase												=  hidden_states.permute(0			,				2			,				1			,				3			,				4     )

						__lowerCAmelCase												=  self.norm(UpperCamelCase     )
						__lowerCAmelCase												=  hidden_states.permute(0			,				3			,				4			,				2			,				1     ).reshape(batch_size * height * width			,				UpperCamelCase			,				UpperCamelCase     )

						__lowerCAmelCase												=  self.proj_in(UpperCamelCase     )

						# 2. Blocks
						for block in self.transformer_blocks:
							__lowerCAmelCase												=  block(
							    UpperCamelCase			,				encoder_hidden_states=UpperCamelCase			,				timestep=UpperCamelCase			,				cross_attention_kwargs=UpperCamelCase			,				class_labels=UpperCamelCase			,				)

						# 3. Output
						__lowerCAmelCase												=  self.proj_out(UpperCamelCase     )
						__lowerCAmelCase												=  (
						    hidden_states[None, None, :]
						    .reshape(UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )
						    .permute(0			,				3			,				4			,				1			,				2     )
						    .contiguous()
						)
						__lowerCAmelCase												=  hidden_states.reshape(UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )

						__lowerCAmelCase												=  hidden_states + residual

						if not return_dict:
							return (output,)

						return TransformerTemporalModelOutput(sample=UpperCamelCase     )