





"""simple docstring"""





import os
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											=						logging.get_logger(__name__)

__A											=						"""▁"""

__A											=						{"""vocab_file""": """sentencepiece.bpe.model"""}

__A											=						{
    """vocab_file""": {
        """xlm-roberta-base""": """https://huggingface.co/xlm-roberta-base/resolve/main/sentencepiece.bpe.model""",
        """xlm-roberta-large""": """https://huggingface.co/xlm-roberta-large/resolve/main/sentencepiece.bpe.model""",
        """xlm-roberta-large-finetuned-conll02-dutch""": (
            """https://huggingface.co/xlm-roberta-large-finetuned-conll02-dutch/resolve/main/sentencepiece.bpe.model"""
        ),
        """xlm-roberta-large-finetuned-conll02-spanish""": (
            """https://huggingface.co/xlm-roberta-large-finetuned-conll02-spanish/resolve/main/sentencepiece.bpe.model"""
        ),
        """xlm-roberta-large-finetuned-conll03-english""": (
            """https://huggingface.co/xlm-roberta-large-finetuned-conll03-english/resolve/main/sentencepiece.bpe.model"""
        ),
        """xlm-roberta-large-finetuned-conll03-german""": (
            """https://huggingface.co/xlm-roberta-large-finetuned-conll03-german/resolve/main/sentencepiece.bpe.model"""
        ),
    }
}

__A											=						{
    """xlm-roberta-base""": 512,
    """xlm-roberta-large""": 512,
    """xlm-roberta-large-finetuned-conll02-dutch""": 512,
    """xlm-roberta-large-finetuned-conll02-spanish""": 512,
    """xlm-roberta-large-finetuned-conll03-english""": 512,
    """xlm-roberta-large-finetuned-conll03-german""": 512,
}






class    _lowerCAmelCase  (	a			):
       """simple docstring"""




       __magic_name__	:List[str]						    =   VOCAB_FILES_NAMES
       __magic_name__	:Any						    =   PRETRAINED_VOCAB_FILES_MAP
       __magic_name__	:int						    =   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       __magic_name__	:Tuple						    =   ["""input_ids""", """attention_mask"""]






       def __init__(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase="<s>"       ,					__UpperCAmelCase="</s>"       ,					__UpperCAmelCase="</s>"       ,					__UpperCAmelCase="<s>"       ,					__UpperCAmelCase="<unk>"       ,					__UpperCAmelCase="<pad>"       ,					__UpperCAmelCase="<mask>"       ,					__UpperCAmelCase = None       ,					**__UpperCAmelCase       ,					):



         '''simple docstring'''

         lowerCAmelCase__     :List[str]											=				AddedToken(__UpperCAmelCase       ,					lstrip=__UpperCAmelCase       ,					rstrip=__UpperCAmelCase ) if isinstance(__UpperCAmelCase       ,					__UpperCAmelCase ) else mask_token

         lowerCAmelCase__     :Dict											=				{} if sp_model_kwargs is None else sp_model_kwargs

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

         lowerCAmelCase__     :Union[str, Any]											=				spm.SentencePieceProcessor(**self.sp_model_kwargs )
         self.sp_model.Load(str(__UpperCAmelCase ) )
         lowerCAmelCase__     :List[str]											=				vocab_file

         # Original fairseq vocab and spm vocab must be "aligned":
         # Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
         # -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
         # fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
         # spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

         # Mimic fairseq token-to-id alignment for the first 4 token
         lowerCAmelCase__     :Union[str, Any]											=				{'<s>': 0, '<pad>': 1, '</s>': 2, '<unk>': 3}

         # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
         lowerCAmelCase__     :Optional[int]											=				1

         lowerCAmelCase__     :Union[str, Any]											=				len(self.sp_model ) + self.fairseq_offset
         lowerCAmelCase__     :List[Any]											=				{v: k for k, v in self.fairseq_tokens_to_ids.items()}






       def __getstate__(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Any											=				self.__dict__.copy()
         lowerCAmelCase__     :Any											=				None
         lowerCAmelCase__     :List[Any]											=				self.sp_model.serialized_model_proto()
         return state






       def __setstate__(       self       ,					__UpperCAmelCase ):



         '''simple docstring'''

         lowerCAmelCase__     :Optional[int]											=				d

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

         lowerCAmelCase__     :Any											=				spm.SentencePieceProcessor(**self.sp_model_kwargs )
         self.sp_model.LoadFromSerializedProto(self.sp_model_proto )






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase = None ):



         '''simple docstring'''

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






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = False ):



         '''simple docstring'''

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

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






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase = None ):



         '''simple docstring'''

         lowerCAmelCase__     :Any											=				[self.sep_token_id]
         lowerCAmelCase__     :Dict											=				[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]






       @property
       def 	snake_case		(       self ):



         '''simple docstring'''

         return len(self.sp_model ) + self.fairseq_offset + 1  # Add the <mask> token






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Optional[int]											=				{self.convert_ids_to_tokens(__UpperCAmelCase ): i for i in range(self.vocab_size )}
         vocab.update(self.added_tokens_encoder )
         return vocab






       def 	snake_case		(       self       ,					__UpperCAmelCase ):



         '''simple docstring'''

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






       def 	snake_case		(       self       ,					__UpperCAmelCase ):



         '''simple docstring'''

         if token in self.fairseq_tokens_to_ids:
           return self.fairseq_tokens_to_ids[token]
         lowerCAmelCase__     :List[str]											=				self.sp_model.PieceToId(__UpperCAmelCase )

         # Need to return unknown token if the SP model returned 0
         return spm_id + self.fairseq_offset if spm_id else self.unk_token_id






       def 	snake_case		(       self       ,					__UpperCAmelCase ):



         '''simple docstring'''

         if index in self.fairseq_ids_to_tokens:
           return self.fairseq_ids_to_tokens[index]
         return self.sp_model.IdToPiece(index - self.fairseq_offset )






       def 	snake_case		(       self       ,					__UpperCAmelCase ):



         '''simple docstring'''

         lowerCAmelCase__     :List[str]											=				''.join(__UpperCAmelCase ).replace(__UpperCAmelCase       ,					' ' ).strip()
         return out_string






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase = None ):



         '''simple docstring'''

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

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

         return (out_vocab_file,)







"""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    _lowerCAmelCase  (	a			):
       """simple docstring"""




       __magic_name__	:torch.FloatTensor






class    _lowerCAmelCase  (	a    ,  a			):
       """simple docstring"""




       @register_to_config
       def __init__(       self       ,					__UpperCAmelCase = 1_6       ,					__UpperCAmelCase = 8_8       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = 1       ,					__UpperCAmelCase = 0.0       ,					__UpperCAmelCase = 3_2       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = False       ,					__UpperCAmelCase = None       ,					__UpperCAmelCase = "geglu"       ,					__UpperCAmelCase = True       ,					__UpperCAmelCase = True       ,					):



         '''simple docstring'''

         super().__init__()
         lowerCAmelCase__     :Dict											=				num_attention_heads
         lowerCAmelCase__     :Any											=				attention_head_dim
         lowerCAmelCase__     :Optional[int]											=				num_attention_heads * attention_head_dim

         lowerCAmelCase__     :Any											=				in_channels

         lowerCAmelCase__     :str											=				torch.nn.GroupNorm(num_groups=__UpperCAmelCase       ,					num_channels=__UpperCAmelCase       ,					eps=1E-6       ,					affine=__UpperCAmelCase )
         lowerCAmelCase__     :int											=				nn.Linear(__UpperCAmelCase       ,					__UpperCAmelCase )

         # 3. Define transformers blocks
         lowerCAmelCase__     :List[Any]											=				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__     :List[Any]											=				nn.Linear(__UpperCAmelCase       ,					__UpperCAmelCase )






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase=None       ,					__UpperCAmelCase=None       ,					__UpperCAmelCase=None       ,					__UpperCAmelCase=1       ,					__UpperCAmelCase=None       ,					__UpperCAmelCase = True       ,					):



         '''simple docstring'''

         lowerCAmelCase__				,							lowerCAmelCase__				,							lowerCAmelCase__				,							lowerCAmelCase__     :int											=				hidden_states.shape
         lowerCAmelCase__     :Tuple											=				batch_frames // num_frames

         lowerCAmelCase__     :str											=				hidden_states

         lowerCAmelCase__     :Union[str, Any]											=				hidden_states[None, :].reshape(__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase )
         lowerCAmelCase__     :str											=				hidden_states.permute(0       ,					2       ,					1       ,					3       ,					4 )

         lowerCAmelCase__     :Optional[int]											=				self.norm(__UpperCAmelCase )
         lowerCAmelCase__     :Optional[Any]											=				hidden_states.permute(0       ,					3       ,					4       ,					2       ,					1 ).reshape(batch_size * height * width       ,					__UpperCAmelCase       ,					__UpperCAmelCase )

         lowerCAmelCase__     :List[Any]											=				self.proj_in(__UpperCAmelCase )

         # 2. Blocks
         for block in self.transformer_blocks:
           lowerCAmelCase__     :Optional[int]											=				block(
               __UpperCAmelCase       ,					encoder_hidden_states=__UpperCAmelCase       ,					timestep=__UpperCAmelCase       ,					cross_attention_kwargs=__UpperCAmelCase       ,					class_labels=__UpperCAmelCase       ,					)

         # 3. Output
         lowerCAmelCase__     :Any											=				self.proj_out(__UpperCAmelCase )
         lowerCAmelCase__     :Dict											=				(
             hidden_states[None, None, :]
             .reshape(__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase )
             .permute(0       ,					3       ,					4       ,					1       ,					2 )
             .contiguous()
         )
         lowerCAmelCase__     :Optional[Any]											=				hidden_states.reshape(__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase )

         lowerCAmelCase__     :List[Any]											=				hidden_states + residual

         if not return_dict:
           return (output,)

         return TransformerTemporalModelOutput(sample=__UpperCAmelCase )
