




"""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_    :
		def __init__(						self				,						UpperCamelCase_				,						UpperCamelCase_=99				,						UpperCamelCase_=13				,						UpperCamelCase_=16				,						UpperCamelCase_=7				,						UpperCamelCase_=True				,						UpperCamelCase_=True				,						UpperCamelCase_=True				,						UpperCamelCase_=False				,						UpperCamelCase_=True				,						UpperCamelCase_=2				,						UpperCamelCase_=32				,						UpperCamelCase_=4				,						UpperCamelCase_=4				,						UpperCamelCase_=30				,						UpperCamelCase_=0				,						UpperCamelCase_=1				,						UpperCamelCase_=2				,						UpperCamelCase_=None				,						)						->       str:
					__lowercase			:  int								=       parent
					__lowercase			:  Tuple								=       batch_size
					__lowercase			:  Optional[Any]								=       decoder_seq_length
					# For common tests
					__lowercase			:  str								=       self.decoder_seq_length
					__lowercase			:  Dict								=       is_training
					__lowercase			:  int								=       use_attention_mask
					__lowercase			:  Dict								=       use_labels

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

					__lowercase			:  List[str]								=       None
					__lowercase			:  Dict								=       decoder_seq_length
					__lowercase			:  Tuple								=       2
					__lowercase			:  Optional[Any]								=       1
		def 	_lowerCamelCase							(						self       )						->       Optional[int]:
					__lowercase			:  Optional[int]								=       ids_tensor([self.batch_size, self.decoder_seq_length]				,						self.vocab_size       )

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

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

					__lowercase			:  str								=       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				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						)						->       str:
					__lowercase			:  Optional[int]								=       True
					__lowercase			:  Tuple								=       TrOCRDecoder(config=A_       ).to(A_       ).eval()
					__lowercase			:  Optional[Any]								=       input_ids[:2]

					input_ids[input_ids == 0] += 1
					# first forward pass
					__lowercase			:  int								=       model(A_				,						use_cache=A_       )
					__lowercase			:  Optional[int]								=       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			:  str								=       outputs['''past_key_values''']

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

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

					__lowercase			:  List[str]								=       model(A_       )['''last_hidden_state''']
					__lowercase			:  Dict								=       model(A_				,						past_key_values=A_       )['''last_hidden_state''']

					# select random slice
					__lowercase			:  str								=       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			:  Union[str, Any]								=       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       )						->       Optional[Any]:
					__lowercase			:  str								=       self.prepare_config_and_inputs()
					__lowercase			,__lowercase			,__lowercase			,__lowercase			:  Dict								=       config_and_inputs

					__lowercase			:  Any								=       {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
					return config, inputs_dict


@require_torch
class   UpperCAmelCase_    (			__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					unittest.TestCase		):
		UpperCamelCase							=(TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
		UpperCamelCase							=(TrOCRForCausalLM,) if is_torch_available() else ()
		UpperCamelCase							={"""text-generation""": TrOCRForCausalLM} if is_torch_available() else {}
		UpperCamelCase							=True
		UpperCamelCase							=False
		def 	_lowerCamelCase							(						self       )						->       Any:
					__lowercase			:  Tuple								=       TrOCRStandaloneDecoderModelTester(self				,						is_training=A_       )
					__lowercase			:  Tuple								=       ConfigTester(self				,						config_class=A_       )
		def 	_lowerCamelCase							(						self       )						->       Optional[int]:
					pass
		def 	_lowerCamelCase							(						self       )						->       List[Any]:
					pass
		def 	_lowerCamelCase							(						self       )						->       Optional[int]:
					pass
		def 	_lowerCamelCase							(						self       )						->       int:
					self.config_tester.run_common_tests()
		def 	_lowerCamelCase							(						self       )						->       List[Any]:
					__lowercase			:  Tuple								=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_decoder_model_past(*A_       )
		def 	_lowerCamelCase							(						self       )						->       Tuple:
					return
		@unittest.skip('''The model doesn\'t support left padding'''       )  # and it's not used enough to be worth fixing :)
		def 	_lowerCamelCase							(						self       )						->       Any:
					pass










import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
_A					:				Optional[int]				=			object()

# For specifying empty leaf dict `{}`
_A					:				Tuple				=			object()


def 							__snake_case							(    lowerCAmelCase_ ,							lowerCAmelCase_	)				->						Dict:
    SCREAMING_SNAKE_CASE__											=			tuple((re.compile(x + '''$'''	) for x in qs)	)
    for i in range(len(lowerCAmelCase_	) - len(lowerCAmelCase_	) + 1	):
        SCREAMING_SNAKE_CASE__											=			[x.match(lowerCAmelCase_	) for x, y in zip(lowerCAmelCase_ ,							ks[i:]	)]
        if matches and all(lowerCAmelCase_	):
            return True
    return False


def 							__snake_case							(    lowerCAmelCase_	)				->						int:
    def replace(lowerCAmelCase_ ,							lowerCAmelCase_	):
        for rule, replacement in rules:
            if _match(lowerCAmelCase_ ,							lowerCAmelCase_	):
                return replacement
        return val

    return replace


def 							__snake_case							(    )				->						Optional[int]:
    return [
        # embeddings
        (("transformer", "wpe", "embedding"), P('''mp''' ,							lowerCAmelCase_	)),
        (("transformer", "wte", "embedding"), P('''mp''' ,							lowerCAmelCase_	)),
        # atention
        (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(lowerCAmelCase_ ,							'''mp'''	)),
        (("attention", "out_proj", "kernel"), P('''mp''' ,							lowerCAmelCase_	)),
        (("attention", "out_proj", "bias"), None),
        # mlp
        (("mlp", "c_fc", "kernel"), P(lowerCAmelCase_ ,							'''mp'''	)),
        (("mlp", "c_fc", "bias"), P('''mp'''	)),
        (("mlp", "c_proj", "kernel"), P('''mp''' ,							lowerCAmelCase_	)),
        (("mlp", "c_proj", "bias"), None),
        # layer norms
        ((r"ln_\d+", "bias"), None),
        ((r"\d+", r"ln_\d+", "scale"), None),
        (("ln_f", "bias"), None),
        (("ln_f", "scale"), None),
    ]



def 							__snake_case							(    lowerCAmelCase_	)				->						Optional[Any]:
    SCREAMING_SNAKE_CASE__											=			_get_partition_rules()
    SCREAMING_SNAKE_CASE__											=			_replacement_rules(lowerCAmelCase_	)
    SCREAMING_SNAKE_CASE__											=			{k: _unmatched for k in flatten_dict(lowerCAmelCase_	)}
    SCREAMING_SNAKE_CASE__											=			{k: replace(lowerCAmelCase_ ,							lowerCAmelCase_	) for k, v in initd.items()}
    assert _unmatched not in result.values(), "Incomplete partition spec."
    return freeze(unflatten_dict(lowerCAmelCase_	)	)
