


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



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

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

		lowerCAmelCase					:    Optional[int]      							= None
		lowerCAmelCase					:    int      							= decoder_seq_length
		lowerCAmelCase					:    int      							= 2
		lowerCAmelCase					:    Optional[int]      							= 1



	def 				lowercase     (					self    ):
		lowerCAmelCase					:    Tuple      							= ids_tensor([self.batch_size, self.decoder_seq_length] ,					self.vocab_size    )

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

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

		lowerCAmelCase					:    Optional[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 				lowercase     (					self ,					snake_case__ ,					snake_case__ ,					snake_case__ ,					snake_case__ ,					):
		lowerCAmelCase					:    List[Any]      							= True
		lowerCAmelCase					:    Union[str, Any]      							= TrOCRDecoder(config=snake_case__    ).to(snake_case__    ).eval()
		lowerCAmelCase					:    List[Any]      							= input_ids[:2]

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

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

		lowerCAmelCase					:    Optional[int]      							= outputs['past_key_values']

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

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

		lowerCAmelCase					:    Any      							= model(snake_case__    )['last_hidden_state']
		lowerCAmelCase					:    Tuple      							= model(snake_case__ ,					past_key_values=snake_case__    )['last_hidden_state']

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

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



	def 				lowercase     (					self    ):
		lowerCAmelCase					:    Union[str, Any]      							= self.prepare_config_and_inputs()
		lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase					:    str      							= config_and_inputs

		lowerCAmelCase					:    List[Any]      							= {'input_ids': input_ids, 'attention_mask': attention_mask}
		return config, inputs_dict



@require_torch
class  lowerCAmelCase				(				a      ,				a      ,				a      ,				unittest.TestCase							):
	_lowerCamelCase   :					List[Any]			      = (TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
	_lowerCamelCase   :					int			      = (TrOCRForCausalLM,) if is_torch_available() else ()
	_lowerCamelCase   :					List[Any]			      = {"""text-generation""": TrOCRForCausalLM} if is_torch_available() else {}
	_lowerCamelCase   :					str			      = True
	_lowerCamelCase   :					int			      = False



	def 				lowercase     (					self    ):
		lowerCAmelCase					:    List[Any]      							= TrOCRStandaloneDecoderModelTester(self ,					is_training=snake_case__    )
		lowerCAmelCase					:    Optional[Any]      							= ConfigTester(self ,					config_class=snake_case__    )



	def 				lowercase     (					self    ):
		pass



	def 				lowercase     (					self    ):
		pass



	def 				lowercase     (					self    ):
		pass



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



	def 				lowercase     (					self    ):
		lowerCAmelCase					:    List[str]      							= self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_decoder_model_past(*snake_case__    )



	def 				lowercase     (					self    ):
		return



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




'''simple docstring'''




import os

import pytest
import yaml

from datasets.features.features import Features, Value
from datasets.info import DatasetInfo, DatasetInfosDict

@pytest.mark.parametrize(
    'files'					,   [
        ['full:README.md', 'dataset_infos.json'],
        ['empty:README.md', 'dataset_infos.json'],
        ['dataset_infos.json'],
        ['full:README.md'],
    ]					,   )
def 			__UpperCamelCase    (    _A    :				str					,   _A    :				List[Any]				)    ->      Union[str, Any]:




	"""simple docstring"""





	lowerCAmelCase					:    Optional[int]      							= tmp_path_factory.mktemp('dset_infos_dir'				)
	if "full:README.md" in files:
		with open(dataset_infos_dir / 'README.md'					,   'w'				) as f:
			f.write('---\ndataset_info:\n  dataset_size: 42\n---'				)
	if "empty:README.md" in files:
		with open(dataset_infos_dir / 'README.md'					,   'w'				) as f:
			f.write(''				)
    # we want to support dataset_infos.json for backward compatibility
	if "dataset_infos.json" in files:
		with open(dataset_infos_dir / 'dataset_infos.json'					,   'w'				) as f:
			f.write('{"default": {"dataset_size": 42}}'				)
	lowerCAmelCase					:    Union[str, Any]      							= DatasetInfosDict.from_directory(_A				)
	assert dataset_infos
	assert dataset_infos["default"].dataset_size == 42

@pytest.mark.parametrize(
    'dataset_info'					,   [
        DatasetInfo(),
        DatasetInfo(
            description='foo'					,   features=Features({'a': Value('int32'				)}				)					,   builder_name='builder'					,   config_name='config'					,   version='1.0.0'					,   splits=[{'name': 'train'}]					,   download_size=42					,   ),
    ]					,   )
def 			__UpperCamelCase    (    _A    :				str					,   _A    :				DatasetInfo				)    ->      Optional[int]:




	"""simple docstring"""





	lowerCAmelCase					:    str      							= str(_A				)
	dataset_info.write_to_directory(_A				)
	lowerCAmelCase					:    List[str]      							= DatasetInfo.from_directory(_A				)
	assert dataset_info == reloaded
	assert os.path.exists(os.path.join(_A					,   'dataset_info.json'				)				)

def 			__UpperCamelCase    (    )    ->      List[str]:




	"""simple docstring"""





	lowerCAmelCase					:    Tuple      							= DatasetInfo(
	    description='foo'					,   citation='bar'					,   homepage='https://foo.bar'					,   license='CC0'					,   features=Features({'a': Value('int32'				)}				)					,   post_processed={}					,   supervised_keys=()					,   task_templates=[]					,   builder_name='builder'					,   config_name='config'					,   version='1.0.0'					,   splits=[{'name': 'train', 'num_examples': 42}]					,   download_checksums={}					,   download_size=13_37					,   post_processing_size=4_42					,   dataset_size=12_34					,   size_in_bytes=13_37 + 4_42 + 12_34					,   )
	lowerCAmelCase					:    Optional[int]      							= dataset_info._to_yaml_dict()
	assert sorted(_A				) == sorted(DatasetInfo._INCLUDED_INFO_IN_YAML				)
	for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
		assert key in dataset_info_yaml_dict
		assert isinstance(dataset_info_yaml_dict[key]					,   (list, dict, int, str)				)
	lowerCAmelCase					:    Any      							= yaml.safe_dump(_A				)
	lowerCAmelCase					:    int      							= yaml.safe_load(_A				)
	assert dataset_info_yaml_dict == reloaded

def 			__UpperCamelCase    (    )    ->      Dict:




	"""simple docstring"""





	lowerCAmelCase					:    Union[str, Any]      							= DatasetInfo()
	lowerCAmelCase					:    List[Any]      							= dataset_info._to_yaml_dict()
	assert dataset_info_yaml_dict == {}






@pytest.mark.parametrize(
    'dataset_infos_dict'					,   [
        DatasetInfosDict(),
        DatasetInfosDict({'default': DatasetInfo()}				),
        DatasetInfosDict({'my_config_name': DatasetInfo()}				),
        DatasetInfosDict(
            {
                'default': DatasetInfo(
                    description='foo'					,   features=Features({'a': Value('int32'				)}				)					,   builder_name='builder'					,   config_name='config'					,   version='1.0.0'					,   splits=[{'name': 'train'}]					,   download_size=42					,   )
            }				),
        DatasetInfosDict(
            {
                'v1': DatasetInfo(dataset_size=42				),
                'v2': DatasetInfo(dataset_size=13_37				),
            }				),
    ]					,   )
def 			__UpperCamelCase    (    _A    :				Tuple					,   _A    :				DatasetInfosDict				)    ->      List[Any]:




	"""simple docstring"""





	lowerCAmelCase					:    Tuple      							= str(_A				)
	dataset_infos_dict.write_to_directory(_A				)
	lowerCAmelCase					:    List[str]      							= DatasetInfosDict.from_directory(_A				)

	# the config_name of the dataset_infos_dict take over the attribute
	for config_name, dataset_info in dataset_infos_dict.items():
		lowerCAmelCase					:    Tuple      							= config_name
		# the yaml representation doesn't include fields like description or citation
		# so we just test that we can recover what we can from the yaml
		lowerCAmelCase					:    Optional[Any]      							= DatasetInfo._from_yaml_dict(dataset_info._to_yaml_dict()				)
	assert dataset_infos_dict == reloaded

	if dataset_infos_dict:
		assert os.path.exists(os.path.join(_A					,   'README.md'				)				)
