


'''simple docstring'''



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






	"""simple docstring"""




	def __init__(       self						,    __a						,    __a=99						,    __a=13						,    __a=16						,    __a=7						,    __a=True						,    __a=True						,    __a=True						,    __a=False						,    __a=True						,    __a=2						,    __a=32						,    __a=4						,    __a=4						,    __a=30						,    __a=0						,    __a=1						,    __a=2						,    __a=None						,    ):

						'''simple docstring'''





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

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

						__a      :	List[str]    =  None
						__a      :	int    =  decoder_seq_length
						__a      :	Dict    =  2
						__a      :	Dict    =  1




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





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

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

						__a      :	int    =  None
						if self.use_labels:
											__a      :	Union[str, Any]    =  ids_tensor([self.batch_size, self.decoder_seq_length]						,    self.vocab_size							)

						__a      :	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 				__UpperCAmelCase							(       self						,    __a						,    __a						,    __a						,    __a						,    ):

						'''simple docstring'''





						__a      :	Optional[Any]    =  True
						__a      :	str    =  TrOCRDecoder(config=__a							).to(__a							).eval()
						__a      :	List[Any]    =  input_ids[:2]

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

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

						__a      :	Tuple    =  outputs['past_key_values']

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

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

						__a      :	Tuple    =  model(__a							)['last_hidden_state']
						__a      :	Dict    =  model(__a						,    past_key_values=__a							)['last_hidden_state']

						# select random slice
						__a      :	Optional[int]    =  ids_tensor((1,)						,    output_from_past.shape[-1]							).item()
						__a      :	Union[str, Any]    =  output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
						__a      :	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 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	str    =  self.prepare_config_and_inputs()
						__a      :	int    =  config_and_inputs

						__a      :	Dict    =  {'input_ids': input_ids, 'attention_mask': attention_mask}
						return config, inputs_dict





@require_torch
class 							__UpperCamelCase    (		lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_		,				unittest.TestCase ):






	"""simple docstring"""




	A_           =   (TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
	A_           =   (TrOCRForCausalLM,) if is_torch_available() else ()
	A_           =   {"text-generation": TrOCRForCausalLM} if is_torch_available() else {}
	A_           =   True
	A_           =   False




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	List[Any]    =  TrOCRStandaloneDecoderModelTester(self						,    is_training=__a							)
						__a      :	str    =  ConfigTester(self						,    config_class=__a							)




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						pass




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						pass




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						pass




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						self.config_tester.run_common_tests()




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						__a      :	int    =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_decoder_model_past(*__a							)




	def 				__UpperCAmelCase							(       self							):

						'''simple docstring'''





						return





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

						'''simple docstring'''





						pass







'''simple docstring'''



from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from torch import nn
from transformers import RobertaPreTrainedModel, XLMRobertaConfig, XLMRobertaModel
from transformers.utils import ModelOutput





@dataclass
class 							__UpperCamelCase    (		lowerCAmelCase_ ):
	A_           =   None
	A_           =   None
	A_           =   None
	A_           =   None





class 							__UpperCamelCase    (		lowerCAmelCase_ ):




	def __init__(       self						,    __a=1						,    __a=0						,    __a=2						,    __a=512						,    __a="cls"						,    __a=False						,    __a=True						,    **__a						,    ):

						'''simple docstring'''





						super().__init__(pad_token_id=__a						,    bos_token_id=__a						,    eos_token_id=__a						,    **__a							)
						__a      :	Any    =  project_dim
						__a      :	Optional[Any]    =  pooler_fn
						__a      :	int    =  learn_encoder
						__a      :	str    =  use_attention_mask





class 							__UpperCamelCase    (		lowerCAmelCase_ ):
	A_           =   [r"pooler", r"logit_scale"]
	A_           =   [r"position_ids", r"predictions.decoder.bias"]
	A_           =   "roberta"
	A_           =   RobertaSeriesConfig




	def __init__(       self						,    __a							):

						'''simple docstring'''





						super().__init__(__a							)
						__a      :	Optional[Any]    =  XLMRobertaModel(__a							)
						__a      :	str    =  nn.Linear(config.hidden_size						,    config.project_dim							)
						__a      :	Optional[int]    =  getattr(__a						,    'has_pre_transformation'						,    __a							)
						if self.has_pre_transformation:
											__a      :	int    =  nn.Linear(config.hidden_size						,    config.project_dim							)
											__a      :	List[str]    =  nn.LayerNorm(config.hidden_size						,    eps=config.layer_norm_eps							)
						self.post_init()





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

						'''simple docstring'''





						__a      :	Optional[Any]    =  return_dict if return_dict is not None else self.config.use_return_dict

						__a      :	Tuple    =  self.base_model(
						    input_ids=__a						,    attention_mask=__a						,    token_type_ids=__a						,    position_ids=__a						,    head_mask=__a						,    inputs_embeds=__a						,    encoder_hidden_states=__a						,    encoder_attention_mask=__a						,    output_attentions=__a						,    output_hidden_states=True if self.has_pre_transformation else output_hidden_states						,    return_dict=__a						,    )

						if self.has_pre_transformation:
											__a      :	Optional[Any]    =  outputs['hidden_states'][-2]
											__a      :	Optional[int]    =  self.pre_LN(__a							)
											__a      :	Union[str, Any]    =  self.transformation_pre(__a							)

											return TransformationModelOutput(
											    projection_state=__a						,    last_hidden_state=outputs.last_hidden_state						,    hidden_states=outputs.hidden_states						,    attentions=outputs.attentions						,    )
						else:
											__a      :	Optional[Any]    =  self.transformation(outputs.last_hidden_state							)
											return TransformationModelOutput(
											    projection_state=__a						,    last_hidden_state=outputs.last_hidden_state						,    hidden_states=outputs.hidden_states						,    attentions=outputs.attentions						,    )



