






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 						lowerCAmelCase   :

					'''simple docstring'''

					def __init__(     self :		List[Any]						,   __a :		Any						,   __a :		Tuple=99						,   __a :		str=13						,   __a :		Tuple=16						,   __a :		Dict=7						,   __a :		Optional[Any]=True						,   __a :		Optional[int]=True						,   __a :		Optional[int]=True						,   __a :		Union[str, Any]=False						,   __a :		Union[str, Any]=True						,   __a :		List[Any]=2						,   __a :		Union[str, Any]=32						,   __a :		Any=4						,   __a :		Dict=4						,   __a :		Optional[int]=30						,   __a :		Optional[Any]=0						,   __a :		Tuple=1						,   __a :		List[str]=2						,   __a :		int=None						,   )	->     List[str]:


												"""simple docstring"""

												__lowercase					:							List[str]			=   parent
												__lowercase					:							Optional[int]			=   batch_size
												__lowercase					:							Optional[Any]			=   decoder_seq_length
												# For common tests
												__lowercase					:							str			=   self.decoder_seq_length
												__lowercase					:							Any			=   is_training
												__lowercase					:							List[Any]			=   use_attention_mask
												__lowercase					:							List[Any]			=   use_labels

												__lowercase					:							Any			=   vocab_size
												__lowercase					:							int			=   d_model
												__lowercase					:							Any			=   d_model
												__lowercase					:							Any			=   decoder_layers
												__lowercase					:							Dict			=   decoder_layers
												__lowercase					:							int			=   decoder_ffn_dim
												__lowercase					:							int			=   decoder_attention_heads
												__lowercase					:							int			=   decoder_attention_heads
												__lowercase					:							List[Any]			=   eos_token_id
												__lowercase					:							str			=   bos_token_id
												__lowercase					:							List[str]			=   pad_token_id
												__lowercase					:							Tuple			=   decoder_start_token_id
												__lowercase					:							List[str]			=   use_cache
												__lowercase					:							List[str]			=   max_position_embeddings

												__lowercase					:							List[str]			=   None
												__lowercase					:							str			=   decoder_seq_length
												__lowercase					:							Union[str, Any]			=   2
												__lowercase					:							Tuple			=   1



					def 					lowerCAmelCase (     self :		List[Any]							)	->     Optional[int]:


												"""simple docstring"""

												__lowercase					:							List[Any]			=   ids_tensor([self.batch_size, self.decoder_seq_length]						,   self.vocab_size							)

												__lowercase					:							Any			=   None
												if self.use_attention_mask:
																			__lowercase					:							Union[str, Any]			=   ids_tensor([self.batch_size, self.decoder_seq_length]						,   vocab_size=2							)

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

												__lowercase					:							Tuple			=   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 					lowerCAmelCase (     self :		int						,   __a :		Optional[int]						,   __a :		Dict						,   __a :		List[Any]						,   __a :		Tuple						,   )	->     List[str]:


												"""simple docstring"""

												__lowercase					:							Union[str, Any]			=   True
												__lowercase					:							Optional[int]			=   TrOCRDecoder(config=__a							).to(__a							).eval()
												__lowercase					:							str			=   input_ids[:2]

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

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

												__lowercase					:							int			=   outputs["""past_key_values"""]

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

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

												__lowercase					:							List[Any]			=   model(__a							)["""last_hidden_state"""]
												__lowercase					:							Union[str, Any]			=   model(__a						,   past_key_values=__a							)["""last_hidden_state"""]

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

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



					def 					lowerCAmelCase (     self :		Any							)	->     List[str]:


												"""simple docstring"""

												__lowercase					:							List[str]			=   self.prepare_config_and_inputs()
												__lowercase	,	__lowercase	,	__lowercase	,	__lowercase					:							Dict			=   config_and_inputs

												__lowercase					:							Optional[Any]			=   {"""input_ids""": input_ids, """attention_mask""": attention_mask}
												return config, inputs_dict







@require_torch
class 						lowerCAmelCase   (       __a					,				__a					,				__a					,				unittest.TestCase	):

					'''simple docstring'''

					_A  :	List[Any]     =					(TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
					_A  :	Optional[Any]     =					(TrOCRForCausalLM,) if is_torch_available() else ()
					_A  :	Dict     =					{'''text-generation''': TrOCRForCausalLM} if is_torch_available() else {}
					_A  :	int     =					True
					_A  :	Union[str, Any]     =					False



					def 					lowerCAmelCase (     self :		List[Any]							)	->     int:


												"""simple docstring"""

												__lowercase					:							Dict			=   TrOCRStandaloneDecoderModelTester(self						,   is_training=__a							)
												__lowercase					:							Tuple			=   ConfigTester(self						,   config_class=__a							)



					def 					lowerCAmelCase (     self :		int							)	->     Optional[Any]:


												"""simple docstring"""

												pass



					def 					lowerCAmelCase (     self :		Union[str, Any]							)	->     Optional[int]:


												"""simple docstring"""

												pass



					def 					lowerCAmelCase (     self :		List[Any]							)	->     int:


												"""simple docstring"""

												pass



					def 					lowerCAmelCase (     self :		Dict							)	->     Tuple:


												"""simple docstring"""

												self.config_tester.run_common_tests()



					def 					lowerCAmelCase (     self :		List[str]							)	->     int:


												"""simple docstring"""

												__lowercase					:							str			=   self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_decoder_model_past(*__a							)



					def 					lowerCAmelCase (     self :		Dict							)	->     Tuple:


												"""simple docstring"""

												return



					@unittest.skip("""The model doesn't support left padding"""							)  # and it's not used enough to be worth fixing :)
					def 					lowerCAmelCase (     self :		Dict							)	->     int:


												"""simple docstring"""

												pass







import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


lowerCamelCase       :       int        =				'''▁'''

lowerCamelCase       :       Optional[int]        =				{'''vocab_file''': '''spiece.model'''}

lowerCamelCase       :       str        =				{
    '''vocab_file''': {'''google/pegasus-xsum''': '''https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model'''}
}

lowerCamelCase       :       str        =				{
    '''google/pegasus-xsum''': 5_12,
}


lowerCamelCase       :       Optional[Any]        =				logging.get_logger(__name__)

class 						lowerCAmelCase   (       __a	):

					'''simple docstring'''

					_A  :	Tuple     =					VOCAB_FILES_NAMES

					_A  :	Tuple     =					VOCAB_FILES_NAMES
					_A  :	Optional[int]     =					PRETRAINED_VOCAB_FILES_MAP
					_A  :	str     =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					_A  :	Optional[int]     =					['''input_ids''', '''attention_mask''']



					def __init__(     self :		Union[str, Any]						,   __a :		int						,   __a :		Any="<pad>"						,   __a :		Optional[int]="</s>"						,   __a :		Union[str, Any]="<unk>"						,   __a :		Optional[int]="<mask_2>"						,   __a :		Optional[int]="<mask_1>"						,   __a :		Dict=None						,   __a :		List[str]=103						,   __a :		Optional[Dict[str, Any]] = None						,   **__a :		List[Any]						,   )	->     None:


												"""simple docstring"""

												__lowercase					:							Tuple			=   offset
												if additional_special_tokens is not None:
																			if not isinstance(__a						,   __a							):
																										raise TypeError(
																										    F"additional_special_tokens should be of type {type(__a							)}, but is"
																										    F" {type(__a							)}"							)

																			__lowercase					:							Dict			=   (
																			    ([mask_token_sent] + additional_special_tokens)
																			    if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
																			    else additional_special_tokens
																			)
																			# fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken
																			additional_special_tokens_extended += [
																			    F"<unk_{i}>" for i in range(len(__a							)						,   self.offset - 1							)
																			]

																			if len(set(__a							)							) != len(__a							):
																										raise ValueError(
																										    """Please make sure that the provided additional_special_tokens do not contain an incorrectly"""
																										    F" shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}."							)
																			__lowercase					:							Optional[Any]			=   additional_special_tokens_extended
												else:
																			__lowercase					:							int			=   [mask_token_sent] if mask_token_sent is not None else []
																			additional_special_tokens += [F"<unk_{i}>" for i in range(2						,   self.offset							)]

												__lowercase					:							str			=   {} if sp_model_kwargs is None else sp_model_kwargs

												super().__init__(
												    eos_token=__a						,   unk_token=__a						,   mask_token=__a						,   pad_token=__a						,   mask_token_sent=__a						,   offset=__a						,   additional_special_tokens=__a						,   sp_model_kwargs=self.sp_model_kwargs						,   **__a						,   )
												__lowercase					:							Optional[Any]			=   mask_token_sent
												__lowercase					:							Dict			=   vocab_file
												__lowercase					:							Optional[Any]			=   spm.SentencePieceProcessor(**self.sp_model_kwargs							)
												self.sp_model.Load(__a							)

												# add special tokens to encoder dict
												__lowercase					:							Dict[int, str]			=   {
												    0: self.pad_token,
												    1: self.eos_token,
												}

												if self.mask_token_sent is not None:
																			self.encoder.update(
																			    {
																			        2: self.mask_token_sent,
																			        3: self.mask_token,
																			    }							)

												if self.offset > 0:
																			# entries 2-104 are only used for pretraining and called <mask_1>, <mask_2>, unk_2, ...unk_102
																			# mask_token_sent is already added to list -> so start at 1
																			self.encoder.update({i + 3: additional_special_tokens[i] for i in range(1						,   self.offset - 1							)}							)

												__lowercase					:							Dict[str, int]			=   {v: k for k, v in self.encoder.items()}



					@property
					def 					lowerCAmelCase (     self :		Dict							)	->     int:


												"""simple docstring"""

												return len(self.sp_model							) + self.offset



					def 					lowerCAmelCase (     self :		int							)	->     Dict[str, int]:


												"""simple docstring"""

												__lowercase					:							Any			=   {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 :		List[str]							)	->     List[str]:


												"""simple docstring"""

												__lowercase					:							Union[str, Any]			=   self.__dict__.copy()
												__lowercase					:							Optional[Any]			=   None
												return state



					def __setstate__(     self :		Tuple						,   __a :		Any							)	->     Tuple:


												"""simple docstring"""

												__lowercase					:							List[str]			=   d

												# for backward compatibility
												if not hasattr(self						,   """sp_model_kwargs"""							):
																			__lowercase					:							List[str]			=   {}

												__lowercase					:							Optional[int]			=   spm.SentencePieceProcessor(**self.sp_model_kwargs							)
												self.sp_model.Load(self.vocab_file							)



					def 					lowerCAmelCase (     self :		Dict						,   __a :		str							)	->     List[str]:


												"""simple docstring"""

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



					def 					lowerCAmelCase (     self :		List[str]						,   __a :		str							)	->     int:


												"""simple docstring"""

												if token in self.decoder:
																			return self.decoder[token]
												elif token in self.added_tokens_decoder:
																			return self.added_tokens_decoder[token]
												__lowercase					:							Union[str, Any]			=   self.sp_model.piece_to_id(__a							)
												return sp_id + self.offset



					def 					lowerCAmelCase (     self :		Dict						,   __a :		int							)	->     str:


												"""simple docstring"""

												if index in self.encoder:
																			return self.encoder[index]
												elif index in self.added_tokens_encoder:
																			return self.added_tokens_encoder[index]
												else:
																			__lowercase					:							List[Any]			=   self.sp_model.IdToPiece(index - self.offset							)
												return token



					def 					lowerCAmelCase (     self :		Union[str, Any]						,   __a :		int							)	->     Dict:


												"""simple docstring"""

												__lowercase					:							Optional[int]			=   []
												__lowercase					:							Tuple			=   """"""
												for token in tokens:
																			# make sure that special tokens are not decoded using sentencepiece model
																			if token in self.all_special_tokens:
																										out_string += self.sp_model.decode(__a							) + token
																										__lowercase					:							str			=   []
																			else:
																										current_sub_tokens.append(__a							)
												out_string += self.sp_model.decode(__a							)
												return out_string.strip()



					def 					lowerCAmelCase (     self :		int						,   __a :		Optional[Any]=False							)	->     int:


												"""simple docstring"""

												return 1



					def 					lowerCAmelCase (     self :		Optional[int]						,   __a :		List[str]							)	->     Optional[Any]:


												"""simple docstring"""

												__lowercase					:							List[Any]			=   set(self.all_special_ids							)  # call it once instead of inside list comp
												all_special_ids.remove(self.unk_token_id							)  # <unk> is only sometimes special

												return [1 if x in all_special_ids else 0 for x in seq]



					def 					lowerCAmelCase (     self :		Union[str, Any]						,   __a :		List						,   __a :		Optional[List] = None						,   __a :		bool = False							)	->     List[int]:


												"""simple docstring"""

												if already_has_special_tokens:
																			return self._special_token_mask(__a							)
												elif token_ids_a is None:
																			return self._special_token_mask(__a							) + [1]
												else:
																			return self._special_token_mask(token_ids_a + token_ids_a							) + [1]



					def 					lowerCAmelCase (     self :		Optional[int]						,   __a :		Dict						,   __a :		Tuple=None							)	->     List[int]:


												"""simple docstring"""

												if token_ids_a is None:
																			return token_ids_a + [self.eos_token_id]
												# We don't expect to process pairs, but leave the pair logic for API consistency
												return token_ids_a + token_ids_a + [self.eos_token_id]



					def 					lowerCAmelCase (     self :		Tuple						,   __a :		str						,   __a :		Optional[str] = None							)	->     Tuple[str]:


												"""simple docstring"""

												if not os.path.isdir(__a							):
																			logger.error(F"Vocabulary path ({save_directory}) should be a directory"							)
																			return
												__lowercase					:							Optional[int]			=   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:
																										__lowercase					:							Any			=   self.sp_model.serialized_model_proto()
																										fi.write(__a							)

												return (out_vocab_file,)