



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



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


           '''simple docstring'''

           __a		 =       parent
           __a		 =       batch_size
           __a		 =       decoder_seq_length
           # For common tests
           __a		 =       self.decoder_seq_length
           __a		 =       is_training
           __a		 =       use_attention_mask
           __a		 =       use_labels

           __a		 =       vocab_size
           __a		 =       d_model
           __a		 =       d_model
           __a		 =       decoder_layers
           __a		 =       decoder_layers
           __a		 =       decoder_ffn_dim
           __a		 =       decoder_attention_heads
           __a		 =       decoder_attention_heads
           __a		 =       eos_token_id
           __a		 =       bos_token_id
           __a		 =       pad_token_id
           __a		 =       decoder_start_token_id
           __a		 =       use_cache
           __a		 =       max_position_embeddings

           __a		 =       None
           __a		 =       decoder_seq_length
           __a		 =       2
           __a		 =       1



       def 	_lowerCamelCase    (	self							:  List[str]):


           '''simple docstring'''

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

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

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

           __a		 =       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							:  Tuple  ,	__SCREAMING_SNAKE_CASE							:  Tuple  ,	__SCREAMING_SNAKE_CASE							:  int  ,	__SCREAMING_SNAKE_CASE							:  Tuple  ,	__SCREAMING_SNAKE_CASE							:  Any  ,	):


           '''simple docstring'''

           __a		 =       True
           __a		 =       TrOCRDecoder(config=__SCREAMING_SNAKE_CASE).to(__SCREAMING_SNAKE_CASE).eval()
           __a		 =       input_ids[:2]

           input_ids[input_ids == 0] += 1
           # first forward pass
           __a		 =       model(__SCREAMING_SNAKE_CASE  ,	use_cache=__SCREAMING_SNAKE_CASE)
           __a		 =       model(__SCREAMING_SNAKE_CASE)
           __a		 =       model(__SCREAMING_SNAKE_CASE  ,	use_cache=__SCREAMING_SNAKE_CASE)

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

           __a		 =       outputs['''past_key_values''']

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

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

           __a		 =       model(__SCREAMING_SNAKE_CASE)['''last_hidden_state''']
           __a		 =       model(__SCREAMING_SNAKE_CASE  ,	past_key_values=__SCREAMING_SNAKE_CASE)['''last_hidden_state''']

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

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





       def 	_lowerCamelCase    (	self							:  List[str]):


           '''simple docstring'''

           __a		 =       self.prepare_config_and_inputs()
           __a ,     __a ,     __a ,     __a		 =       config_and_inputs

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


@require_torch
class 					_A (			__UpperCAmelCase   ,__UpperCAmelCase   ,__UpperCAmelCase   ,unittest.TestCase					):
       UpperCamelCase__						:					Optional[int]		      =		(TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
       UpperCamelCase__						:					Optional[Any]		      =		(TrOCRForCausalLM,) if is_torch_available() else ()
       UpperCamelCase__						:					int		      =		{'''text-generation''': TrOCRForCausalLM} if is_torch_available() else {}
       UpperCamelCase__						:					Optional[Any]		      =		True
       UpperCamelCase__						:					int		      =		False



       def 	_lowerCamelCase    (	self							:  List[str]):


           '''simple docstring'''

           __a		 =       TrOCRStandaloneDecoderModelTester(self  ,	is_training=__SCREAMING_SNAKE_CASE)
           __a		 =       ConfigTester(self  ,	config_class=__SCREAMING_SNAKE_CASE)



       def 	_lowerCamelCase    (	self							:  List[Any]):


           '''simple docstring'''

           pass



       def 	_lowerCamelCase    (	self							:  Optional[Any]):


           '''simple docstring'''

           pass



       def 	_lowerCamelCase    (	self							:  Dict):


           '''simple docstring'''

           pass



       def 	_lowerCamelCase    (	self							:  str):


           '''simple docstring'''

           self.config_tester.run_common_tests()



       def 	_lowerCamelCase    (	self							:  List[Any]):


           '''simple docstring'''

           __a		 =       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_decoder_model_past(*__SCREAMING_SNAKE_CASE)



       def 	_lowerCamelCase    (	self							:  List[Any]):


           '''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							:  Any):


           '''simple docstring'''

           pass










from unittest import TestCase

from datasets import Sequence, Value
from datasets.arrow_dataset import Dataset

class 					_A (			__UpperCAmelCase					):



       def 	_lowerCamelCase    (	self							:  int):


           '''simple docstring'''

           return [
               {"col_1": 3, "col_2": "a"},
               {"col_1": 2, "col_2": "b"},
               {"col_1": 1, "col_2": "c"},
               {"col_1": 0, "col_2": "d"},
           ]



       def 	_lowerCamelCase    (	self							:  Tuple):


           '''simple docstring'''

           __a		 =       {'''col_1''': [3, 2, 1, 0], '''col_2''': ['''a''', '''b''', '''c''', '''d''']}
           return Dataset.from_dict(__SCREAMING_SNAKE_CASE)



       def 	_lowerCamelCase    (	self							:  Any):


           '''simple docstring'''

           __a		 =       self._create_example_records()
           __a		 =       Dataset.from_list(__SCREAMING_SNAKE_CASE)
           self.assertListEqual(dset.column_names  ,	['''col_1''', '''col_2'''])
           for i, r in enumerate(__SCREAMING_SNAKE_CASE):
               self.assertDictEqual(__SCREAMING_SNAKE_CASE  ,	example_records[i])



       def 	_lowerCamelCase    (	self							:  Optional[Any]):


           '''simple docstring'''

           __a		 =       self._create_example_records()
           __a		 =       Dataset.from_list(__SCREAMING_SNAKE_CASE)
           __a		 =       Dataset.from_dict({k: [r[k] for r in example_records] for k in example_records[0]})
           self.assertEqual(dset.info  ,	dset_from_dict.info)



       def 	_lowerCamelCase    (	self							:  int):  # checks what happens with missing columns


           '''simple docstring'''

           __a		 =       [{'''col_1''': 1}, {'''col_2''': '''x'''}]
           __a		 =       Dataset.from_list(__SCREAMING_SNAKE_CASE)
           self.assertDictEqual(dset[0]  ,	{'''col_1''': 1})
           self.assertDictEqual(dset[1]  ,	{'''col_1''': None})  # NB: first record is used for columns



       def 	_lowerCamelCase    (	self							:  Optional[Any]):  # checks if the type can be inferred from the second record


           '''simple docstring'''

           __a		 =       [{'''col_1''': []}, {'''col_1''': [1, 2]}]
           __a		 =       Dataset.from_list(__SCREAMING_SNAKE_CASE)
           self.assertEqual(dset.info.features['''col_1''']  ,	Sequence(Value('''int64''')))





       def 	_lowerCamelCase    (	self							:  List[Any]):


           '''simple docstring'''

           __a		 =       Dataset.from_list([])
           self.assertEqual(len(__SCREAMING_SNAKE_CASE)  ,	0)
           self.assertListEqual(dset.column_names  ,	[])





