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




							def __init__(    self       ,A__       ,A__=9_9       ,A__=1_3       ,A__=1_6       ,A__=7       ,A__=True       ,A__=True       ,A__=True       ,A__=False       ,A__=True       ,A__=2       ,A__=3_2       ,A__=4       ,A__=4       ,A__=3_0       ,A__=0       ,A__=1       ,A__=2       ,A__=None       ,):
												lowercase						     =						parent
												lowercase						     =						batch_size
												lowercase						     =						decoder_seq_length
												# For common tests
												lowercase						     =						self.decoder_seq_length
												lowercase						     =						is_training
												lowercase						     =						use_attention_mask
												lowercase						     =						use_labels

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

												lowercase						     =						None
												lowercase						     =						decoder_seq_length
												lowercase						     =						2
												lowercase						     =						1




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

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

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

												lowercase						     =						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 			A__				(    self       ,A__       ,A__       ,A__       ,A__       ,):
												lowercase						     =						True
												lowercase						     =						TrOCRDecoder(config=A__).to(A__).eval()
												lowercase						     =						input_ids[:2]

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

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

												lowercase						     =						outputs['''past_key_values''']

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

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

												lowercase						     =						model(A__)['''last_hidden_state''']
												lowercase						     =						model(A__       ,past_key_values=A__)['''last_hidden_state''']

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

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





							def 			A__				(    self):
												lowercase						     =						self.prepare_config_and_inputs()
												lowercase     ,						lowercase     ,						lowercase     ,						lowercase						     =						config_and_inputs

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







@require_torch
class     lowercase		( SCREAMING_SNAKE_CASE__	,  SCREAMING_SNAKE_CASE__	,  SCREAMING_SNAKE_CASE__	,  unittest.TestCase		):
							lowercase_	:	Any	 =(TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
							lowercase_	:	Dict	 =(TrOCRForCausalLM,) if is_torch_available() else ()
							lowercase_	:	int	 ={'''text-generation''': TrOCRForCausalLM} if is_torch_available() else {}
							lowercase_	:	List[Any]	 =True
							lowercase_	:	int	 =False




							def 			A__				(    self):
												lowercase						     =						TrOCRStandaloneDecoderModelTester(self       ,is_training=A__)
												lowercase						     =						ConfigTester(self       ,config_class=A__)




							def 			A__				(    self):
												pass




							def 			A__				(    self):
												pass




							def 			A__				(    self):
												pass




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




							def 			A__				(    self):
												lowercase						     =						self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_decoder_model_past(*A__)




							def 			A__				(    self):
												return





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






import unittest

from transformers import AlbertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import (
							    MODEL_FOR_PRETRAINING_MAPPING,
							    AlbertForMaskedLM,
							    AlbertForMultipleChoice,
							    AlbertForPreTraining,
							    AlbertForQuestionAnswering,
							    AlbertForSequenceClassification,
							    AlbertForTokenClassification,
							    AlbertModel,
							)
							from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST







class     lowercase		:




							def __init__(    self       ,A__       ,A__=1_3       ,A__=7       ,A__=True       ,A__=True       ,A__=True       ,A__=True       ,A__=9_9       ,A__=1_6       ,A__=3_6       ,A__=6       ,A__=6       ,A__=6       ,A__=3_7       ,A__="gelu"       ,A__=0.1       ,A__=0.1       ,A__=5_1_2       ,A__=1_6       ,A__=2       ,A__=0.02       ,A__=3       ,A__=4       ,A__=None       ,):
												lowercase						     =						parent
												lowercase						     =						batch_size
												lowercase						     =						seq_length
												lowercase						     =						is_training
												lowercase						     =						use_input_mask
												lowercase						     =						use_token_type_ids
												lowercase						     =						use_labels
												lowercase						     =						vocab_size
												lowercase						     =						embedding_size
												lowercase						     =						hidden_size
												lowercase						     =						num_hidden_layers
												lowercase						     =						num_hidden_groups
												lowercase						     =						num_attention_heads
												lowercase						     =						intermediate_size
												lowercase						     =						hidden_act
												lowercase						     =						hidden_dropout_prob
												lowercase						     =						attention_probs_dropout_prob
												lowercase						     =						max_position_embeddings
												lowercase						     =						type_vocab_size
												lowercase						     =						type_sequence_label_size
												lowercase						     =						initializer_range
												lowercase						     =						num_labels
												lowercase						     =						num_choices
												lowercase						     =						scope




							def 			A__				(    self):
												lowercase						     =						ids_tensor([self.batch_size, self.seq_length]       ,self.vocab_size)

												lowercase						     =						None
												if self.use_input_mask:
																	lowercase						     =						random_attention_mask([self.batch_size, self.seq_length])

												lowercase						     =						None
												if self.use_token_type_ids:
																	lowercase						     =						ids_tensor([self.batch_size, self.seq_length]       ,self.type_vocab_size)

												lowercase						     =						None
												lowercase						     =						None
												lowercase						     =						None
												if self.use_labels:
																	lowercase						     =						ids_tensor([self.batch_size]       ,self.type_sequence_label_size)
																	lowercase						     =						ids_tensor([self.batch_size, self.seq_length]       ,self.num_labels)
																	lowercase						     =						ids_tensor([self.batch_size]       ,self.num_choices)

												lowercase						     =						self.get_config()

												return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




							def 			A__				(    self):
												return AlbertConfig(
												    vocab_size=self.vocab_size       ,hidden_size=self.hidden_size       ,num_hidden_layers=self.num_hidden_layers       ,num_attention_heads=self.num_attention_heads       ,intermediate_size=self.intermediate_size       ,hidden_act=self.hidden_act       ,hidden_dropout_prob=self.hidden_dropout_prob       ,attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,max_position_embeddings=self.max_position_embeddings       ,type_vocab_size=self.type_vocab_size       ,initializer_range=self.initializer_range       ,num_hidden_groups=self.num_hidden_groups       ,)




							def 			A__				(    self       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__):
												lowercase						     =						AlbertModel(config=A__)
												model.to(A__)
												model.eval()
												lowercase						     =						model(A__       ,attention_mask=A__       ,token_type_ids=A__)
												lowercase						     =						model(A__       ,token_type_ids=A__)
												lowercase						     =						model(A__)
												self.parent.assertEqual(result.last_hidden_state.shape       ,(self.batch_size, self.seq_length, self.hidden_size))
												self.parent.assertEqual(result.pooler_output.shape       ,(self.batch_size, self.hidden_size))




							def 			A__				(    self       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__):
												lowercase						     =						AlbertForPreTraining(config=A__)
												model.to(A__)
												model.eval()
												lowercase						     =						model(
												    A__       ,attention_mask=A__       ,token_type_ids=A__       ,labels=A__       ,sentence_order_label=A__       ,)
												self.parent.assertEqual(result.prediction_logits.shape       ,(self.batch_size, self.seq_length, self.vocab_size))
												self.parent.assertEqual(result.sop_logits.shape       ,(self.batch_size, config.num_labels))




							def 			A__				(    self       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__):
												lowercase						     =						AlbertForMaskedLM(config=A__)
												model.to(A__)
												model.eval()
												lowercase						     =						model(A__       ,attention_mask=A__       ,token_type_ids=A__       ,labels=A__)
												self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.seq_length, self.vocab_size))




							def 			A__				(    self       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__):
												lowercase						     =						AlbertForQuestionAnswering(config=A__)
												model.to(A__)
												model.eval()
												lowercase						     =						model(
												    A__       ,attention_mask=A__       ,token_type_ids=A__       ,start_positions=A__       ,end_positions=A__       ,)
												self.parent.assertEqual(result.start_logits.shape       ,(self.batch_size, self.seq_length))
												self.parent.assertEqual(result.end_logits.shape       ,(self.batch_size, self.seq_length))




							def 			A__				(    self       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__):
												lowercase						     =						self.num_labels
												lowercase						     =						AlbertForSequenceClassification(A__)
												model.to(A__)
												model.eval()
												lowercase						     =						model(A__       ,attention_mask=A__       ,token_type_ids=A__       ,labels=A__)
												self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.num_labels))




							def 			A__				(    self       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__):
												lowercase						     =						self.num_labels
												lowercase						     =						AlbertForTokenClassification(config=A__)
												model.to(A__)
												model.eval()
												lowercase						     =						model(A__       ,attention_mask=A__       ,token_type_ids=A__       ,labels=A__)
												self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.seq_length, self.num_labels))




							def 			A__				(    self       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__       ,A__):
												lowercase						     =						self.num_choices
												lowercase						     =						AlbertForMultipleChoice(config=A__)
												model.to(A__)
												model.eval()
												lowercase						     =						input_ids.unsqueeze(1).expand(-1       ,self.num_choices       ,-1).contiguous()
												lowercase						     =						token_type_ids.unsqueeze(1).expand(-1       ,self.num_choices       ,-1).contiguous()
												lowercase						     =						input_mask.unsqueeze(1).expand(-1       ,self.num_choices       ,-1).contiguous()
												lowercase						     =						model(
												    A__       ,attention_mask=A__       ,token_type_ids=A__       ,labels=A__       ,)
												self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.num_choices))





							def 			A__				(    self):
												lowercase						     =						self.prepare_config_and_inputs()
												(
												    (
												    lowercase
												)     ,						(
												    lowercase
												)     ,						(
												    lowercase
												)     ,						(
												    lowercase
												)     ,						(
												    lowercase
												)     ,						(
												    lowercase
												)     ,						(
												    lowercase
												)     ,						
												)						     =						config_and_inputs
												lowercase						     =						{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
												return config, inputs_dict







@require_torch
class     lowercase		( SCREAMING_SNAKE_CASE__	,  SCREAMING_SNAKE_CASE__	,  unittest.TestCase		):
							lowercase_	:	Union[str, Any]	 =(
							    (
							        AlbertModel,
							        AlbertForPreTraining,
							        AlbertForMaskedLM,
							        AlbertForMultipleChoice,
							        AlbertForSequenceClassification,
							        AlbertForTokenClassification,
							        AlbertForQuestionAnswering,
							    )
							    if is_torch_available()
							    else ()
							)
							lowercase_	:	int	 =(
							    {
							        '''feature-extraction''': AlbertModel,
							        '''fill-mask''': AlbertForMaskedLM,
							        '''question-answering''': AlbertForQuestionAnswering,
							        '''text-classification''': AlbertForSequenceClassification,
							        '''token-classification''': AlbertForTokenClassification,
							        '''zero-shot''': AlbertForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							lowercase_	:	str	 =True




							def 			A__				(    self       ,A__       ,A__       ,A__=False):
												lowercase						     =						super()._prepare_for_class(A__       ,A__       ,return_labels=A__)

												if return_labels:
																	if model_class in get_values(A__):
																						lowercase						     =						torch.zeros(
																						    (self.model_tester.batch_size, self.model_tester.seq_length)       ,dtype=torch.long       ,device=A__)
																						lowercase						     =						torch.zeros(
																						    self.model_tester.batch_size       ,dtype=torch.long       ,device=A__)
												return inputs_dict




							def 			A__				(    self):
												lowercase						     =						AlbertModelTester(self)
												lowercase						     =						ConfigTester(self       ,config_class=A__       ,hidden_size=3_7)




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




							def 			A__				(    self):
												lowercase						     =						self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*A__)




							def 			A__				(    self):
												lowercase						     =						self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_pretraining(*A__)




							def 			A__				(    self):
												lowercase						     =						self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_masked_lm(*A__)




							def 			A__				(    self):
												lowercase						     =						self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_multiple_choice(*A__)




							def 			A__				(    self):
												lowercase						     =						self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_question_answering(*A__)




							def 			A__				(    self):
												lowercase						     =						self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_sequence_classification(*A__)




							def 			A__				(    self):
												lowercase						     =						self.model_tester.prepare_config_and_inputs()
												for type in ["absolute", "relative_key", "relative_key_query"]:
																	lowercase						     =						type
																	self.model_tester.create_and_check_model(*A__)





							@slow
							def 			A__				(    self):
												for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	lowercase						     =						AlbertModel.from_pretrained(A__)
																	self.assertIsNotNone(A__)







@require_torch
class     lowercase		( unittest.TestCase		):




							@slow
							def 			A__				(    self):
												lowercase						     =						AlbertModel.from_pretrained('''albert-base-v2''')
												lowercase						     =						torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]])
												lowercase						     =						torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
												with torch.no_grad():
																	lowercase						     =						model(A__       ,attention_mask=A__)[0]
												lowercase						     =						torch.Size((1, 1_1, 7_6_8))
												self.assertEqual(output.shape       ,A__)
												lowercase						     =						torch.tensor(
												    [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]])

												self.assertTrue(torch.allclose(output[:, 1:4, 1:4]       ,A__       ,atol=1E-4))




