


import unittest

from transformers import TrOCRConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers.models.trocr.modeling_trocr import TrOCRDecoder, TrOCRForCausalLM







@require_torch
class 		_snake_case  :




 def __init__(		self					,  _lowerCamelCase					,  _lowerCamelCase=99					,  _lowerCamelCase=13					,  _lowerCamelCase=16					,  _lowerCamelCase=7					,  _lowerCamelCase=True					,  _lowerCamelCase=True					,  _lowerCamelCase=True					,  _lowerCamelCase=False					,  _lowerCamelCase=True					,  _lowerCamelCase=2					,  _lowerCamelCase=32					,  _lowerCamelCase=4					,  _lowerCamelCase=4					,  _lowerCamelCase=30					,  _lowerCamelCase=0					,  _lowerCamelCase=1					,  _lowerCamelCase=2					,  _lowerCamelCase=None					,  ):
       a						:Union[str, Any] 							=       parent
       a						:Dict 							=       batch_size
       a						:Tuple 							=       decoder_seq_length
       # For common tests
       a						:int 							=       self.decoder_seq_length
       a						:Optional[int] 							=       is_training
       a						:Optional[Any] 							=       use_attention_mask
       a						:Tuple 							=       use_labels

       a						:Any 							=       vocab_size
       a						:Union[str, Any] 							=       d_model
       a						:str 							=       d_model
       a						:int 							=       decoder_layers
       a						:Tuple 							=       decoder_layers
       a						:Optional[int] 							=       decoder_ffn_dim
       a						:str 							=       decoder_attention_heads
       a						:Optional[int] 							=       decoder_attention_heads
       a						:List[Any] 							=       eos_token_id
       a						:Tuple 							=       bos_token_id
       a						:Any 							=       pad_token_id
       a						:Union[str, Any] 							=       decoder_start_token_id
       a						:Optional[Any] 							=       use_cache
       a						:Optional[Any] 							=       max_position_embeddings

       a						:Dict 							=       None
       a						:Optional[Any] 							=       decoder_seq_length
       a						:Optional[int] 							=       2
       a						:int 							=       1




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       a						:Any 							=       ids_tensor([self.batch_size, self.decoder_seq_length]					,  self.vocab_size    )

       a						:int 							=       None
       if self.use_attention_mask:
             a						:Optional[int] 							=       ids_tensor([self.batch_size, self.decoder_seq_length]					,  vocab_size=2    )

       a						:Dict 							=       None
       if self.use_labels:
             a						:Any 							=       ids_tensor([self.batch_size, self.decoder_seq_length]					,  self.vocab_size    )

       a						:List[Any] 							=       TrOCRConfig(
           vocab_size=self.vocab_size					,  d_model=self.d_model					,  decoder_layers=self.decoder_layers					,  decoder_ffn_dim=self.decoder_ffn_dim					,  decoder_attention_heads=self.decoder_attention_heads					,  eos_token_id=self.eos_token_id					,  bos_token_id=self.bos_token_id					,  use_cache=self.use_cache					,  pad_token_id=self.pad_token_id					,  decoder_start_token_id=self.decoder_start_token_id					,  max_position_embeddings=self.max_position_embeddings					,  )

       return (config, input_ids, attention_mask, lm_labels)




 def        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase					,  ):
       a						:int 							=       True
       a						:Any 							=       TrOCRDecoder(config=_lowerCamelCase    ).to(_lowerCamelCase    ).eval()
       a						:Tuple 							=       input_ids[:2]

       input_ids[input_ids == 0] += 1
       # first forward pass
       a						:Any 							=       model(_lowerCamelCase					,  use_cache=_lowerCamelCase    )
       a						:List[Any] 							=       model(_lowerCamelCase    )
       a						:Union[str, Any] 							=       model(_lowerCamelCase					,  use_cache=_lowerCamelCase    )

       self.parent.assertTrue(len(_lowerCamelCase    ) == len(_lowerCamelCase    )    )
       self.parent.assertTrue(len(_lowerCamelCase    ) == len(_lowerCamelCase    ) + 1    )

       a						:Optional[Any] 							=       outputs['''past_key_values''']

       # create hypothetical next token and extent to next_input_ids
       a						:int 							=       ids_tensor((2, 1)					,  config.vocab_size - 1    ) + 1

       # append to next input_ids and
       a						:str 							=       torch.cat([input_ids, next_tokens]					,  dim=-1    )

       a						:List[Any] 							=       model(_lowerCamelCase    )['''last_hidden_state''']
       a						:List[Any] 							=       model(_lowerCamelCase					,  past_key_values=_lowerCamelCase    )['''last_hidden_state''']

       # select random slice
       a						:List[str] 							=       ids_tensor((1,)					,  output_from_past.shape[-1]    ).item()
       a						:Optional[int] 							=       output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
       a						:List[Any] 							=       output_from_past[:, 0, random_slice_idx].detach()

       # test that outputs are equal for slice
       assert torch.allclose(_lowerCamelCase					,  _lowerCamelCase					,  atol=1e-3    )




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       a						:Dict 							=       self.prepare_config_and_inputs()
       a							,				a							,				a							,				a						:Union[str, Any] 							=       config_and_inputs

       a						:Optional[Any] 							=       {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
       return config, inputs_dict







@require_torch
class 		_snake_case  (      _snake_case						,    _snake_case						,    _snake_case						,    unittest.TestCase							):
 SCREAMING_SNAKE_CASE__      =       (TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
 SCREAMING_SNAKE_CASE__      =       (TrOCRForCausalLM,) if is_torch_available() else ()
 SCREAMING_SNAKE_CASE__      =       {'text-generation': TrOCRForCausalLM} if is_torch_available() else {}
 SCREAMING_SNAKE_CASE__      =       True
 SCREAMING_SNAKE_CASE__      =       False




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       a						:Optional[int] 							=       TrOCRStandaloneDecoderModelTester(self					,  is_training=_lowerCamelCase    )
       a						:Tuple 							=       ConfigTester(self					,  config_class=_lowerCamelCase    )




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       pass




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       pass




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       pass




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       self.config_tester.run_common_tests()




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       a						:str 							=       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_decoder_model_past(*_lowerCamelCase    )




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       return




 @unittest.skip('''The model doesn\'t support left padding'''    )  # and it's not used enough to be worth fixing :)
 def        SCREAMING_SNAKE_CASE__				(		self    ):
       pass









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


snake_case				:    Dict       				=     logging.get_logger(__name__)

snake_case				:    Tuple       				=     '''▁'''

snake_case				:    Any       				=     {'''vocab_file''': '''sentencepiece.bpe.model'''}

snake_case				:    Tuple       				=     {
    '''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'''
        ),
    }
}

snake_case				:    int       				=     {
    '''xlm-roberta-base''': 5_12,
    '''xlm-roberta-large''': 5_12,
    '''xlm-roberta-large-finetuned-conll02-dutch''': 5_12,
    '''xlm-roberta-large-finetuned-conll02-spanish''': 5_12,
    '''xlm-roberta-large-finetuned-conll03-english''': 5_12,
    '''xlm-roberta-large-finetuned-conll03-german''': 5_12,
}







class 		_snake_case  (      _snake_case							):
 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']




 def __init__(		self					,  _lowerCamelCase					,  _lowerCamelCase="<s>"					,  _lowerCamelCase="</s>"					,  _lowerCamelCase="</s>"					,  _lowerCamelCase="<s>"					,  _lowerCamelCase="<unk>"					,  _lowerCamelCase="<pad>"					,  _lowerCamelCase="<mask>"					,  _lowerCamelCase = None					,  **_lowerCamelCase					,  ):
       # Mask token behave like a normal word, i.e. include the space before it
       a						:Optional[int] 							=       AddedToken(_lowerCamelCase					,  lstrip=_lowerCamelCase					,  rstrip=_lowerCamelCase    ) if isinstance(_lowerCamelCase					,  _lowerCamelCase    ) else mask_token

       a						:int 							=       {} if sp_model_kwargs is None else sp_model_kwargs

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

       a						:Optional[int] 							=       spm.SentencePieceProcessor(**self.sp_model_kwargs    )
       self.sp_model.Load(str(_lowerCamelCase    )    )
       a						: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
       a						:Tuple 							=       {'''<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
       a						:List[str] 							=       1

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




 def __getstate__(		self    ):
       a						:List[str] 							=       self.__dict__.copy()
       a						:Optional[int] 							=       None
       a						:int 							=       self.sp_model.serialized_model_proto()
       return state




 def __setstate__(		self					,  _lowerCamelCase    ):
       a						:Union[str, Any] 							=       d

       # for backward compatibility
       if not hasattr(self					,  '''sp_model_kwargs'''    ):
             a						:Union[str, Any] 							=       {}

       a						:Optional[int] 							=       spm.SentencePieceProcessor(**self.sp_model_kwargs    )
       self.sp_model.LoadFromSerializedProto(self.sp_model_proto    )




 def        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase					,  _lowerCamelCase = None    ):

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




 def        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase					,  _lowerCamelCase = None					,  _lowerCamelCase = False    ):

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

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




 def        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase					,  _lowerCamelCase = None    ):
       a						:int 							=       [self.sep_token_id]
       a						:int 							=       [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        SCREAMING_SNAKE_CASE__				(		self    ):
       return len(self.sp_model    ) + self.fairseq_offset + 1  # Add the <mask> token




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       a						:Any 							=       {self.convert_ids_to_tokens(_lowerCamelCase    ): i for i in range(self.vocab_size    )}
       vocab.update(self.added_tokens_encoder    )
       return vocab




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




 def        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase    ):
       if token in self.fairseq_tokens_to_ids:
             return self.fairseq_tokens_to_ids[token]
       a						:Optional[Any] 							=       self.sp_model.PieceToId(_lowerCamelCase    )

       # 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        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase    ):
       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        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase    ):
       a						:Tuple 							=       ''''''.join(_lowerCamelCase    ).replace(_lowerCamelCase					,  ''' '''    ).strip()
       return out_string




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

       if os.path.abspath(self.vocab_file    ) != os.path.abspath(_lowerCamelCase    ) and os.path.isfile(self.vocab_file    ):
             copyfile(self.vocab_file					,  _lowerCamelCase    )
       elif not os.path.isfile(self.vocab_file    ):
             with open(_lowerCamelCase					,  '''wb'''    ) as fi:
                   a						:List[Any] 							=       self.sp_model.serialized_model_proto()
                   fi.write(_lowerCamelCase    )

       return (out_vocab_file,)





