
import unittest

from transformers import ErnieMConfig, is_torch_available
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 (
			    ErnieMForInformationExtraction,
			    ErnieMForMultipleChoice,
			    ErnieMForQuestionAnswering,
			    ErnieMForSequenceClassification,
			    ErnieMForTokenClassification,
			    ErnieMModel,
			)
			from transformers.models.ernie_m.modeling_ernie_m import ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST





class      _lowerCAmelCase :




				def __init__( self    ,    _UpperCamelCase    ,    _UpperCamelCase=13    ,    _UpperCamelCase=7    ,    _UpperCamelCase=True    ,    _UpperCamelCase=True    ,    _UpperCamelCase=True    ,    _UpperCamelCase=99    ,    _UpperCamelCase=32    ,    _UpperCamelCase=5    ,    _UpperCamelCase=4    ,    _UpperCamelCase=37    ,    _UpperCamelCase="gelu"    ,    _UpperCamelCase=0.1    ,    _UpperCamelCase=0.1    ,    _UpperCamelCase=512    ,    _UpperCamelCase=16    ,    _UpperCamelCase=2    ,    _UpperCamelCase=0.02    ,    _UpperCamelCase=3    ,    _UpperCamelCase=4    ,    _UpperCamelCase=None    ,    )    ->				List[str]:
							lowerCAmelCase_      =     parent
							lowerCAmelCase_      =     batch_size
							lowerCAmelCase_      =     seq_length
							lowerCAmelCase_      =     is_training
							lowerCAmelCase_      =     use_input_mask
							lowerCAmelCase_      =     use_labels
							lowerCAmelCase_      =     vocab_size
							lowerCAmelCase_      =     hidden_size
							lowerCAmelCase_      =     num_hidden_layers
							lowerCAmelCase_      =     num_attention_heads
							lowerCAmelCase_      =     intermediate_size
							lowerCAmelCase_      =     hidden_act
							lowerCAmelCase_      =     hidden_dropout_prob
							lowerCAmelCase_      =     attention_probs_dropout_prob
							lowerCAmelCase_      =     max_position_embeddings
							lowerCAmelCase_      =     type_vocab_size
							lowerCAmelCase_      =     type_sequence_label_size
							lowerCAmelCase_      =     initializer_range
							lowerCAmelCase_      =     num_labels
							lowerCAmelCase_      =     num_choices
							lowerCAmelCase_      =     scope




				def 					__a						( self					)    ->				List[Any]:
							lowerCAmelCase_      =     ids_tensor([self.batch_size, self.seq_length]    ,    self.vocab_size					)

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

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

							lowerCAmelCase_      =     self.get_config()

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




				def 					__a						( self					)    ->				str:
							lowerCAmelCase_      =     ids_tensor([self.batch_size, self.seq_length]    ,    self.vocab_size					)

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

							return config, input_ids, input_mask




				def 					__a						( self					)    ->				Optional[int]:
							return ErnieMConfig(
							    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    ,    )




				def 					__a						( self    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase					)    ->				Optional[Any]:
							lowerCAmelCase_      =     ErnieMModel(config=UpperCamelCase__					)
							model.to(UpperCamelCase__					)
							model.eval()
							lowerCAmelCase_      =     model(UpperCamelCase__    ,    return_dict=UpperCamelCase__					)
							self.parent.assertEqual(result.last_hidden_state.shape    ,    (self.batch_size, self.seq_length, self.hidden_size)					)




				def 					__a						( self    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase					)    ->				Dict:
							lowerCAmelCase_      =     ErnieMForQuestionAnswering(config=UpperCamelCase__					)
							model.to(UpperCamelCase__					)
							model.eval()
							lowerCAmelCase_      =     model(
							    UpperCamelCase__    ,    attention_mask=UpperCamelCase__    ,    start_positions=UpperCamelCase__    ,    end_positions=UpperCamelCase__    ,    )
							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    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase					)    ->				Optional[Any]:
							lowerCAmelCase_      =     ErnieMForInformationExtraction(config=UpperCamelCase__					)
							model.to(UpperCamelCase__					)
							model.eval()
							lowerCAmelCase_      =     torch.ones_like(UpperCamelCase__    ,    dtype=torch.floataa					)
							lowerCAmelCase_      =     model(
							    UpperCamelCase__    ,    attention_mask=UpperCamelCase__    ,    start_positions=UpperCamelCase__    ,    end_positions=UpperCamelCase__    ,    )
							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    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase					)    ->				List[Any]:
							lowerCAmelCase_      =     self.num_labels
							lowerCAmelCase_      =     ErnieMForSequenceClassification(UpperCamelCase__					)
							model.to(UpperCamelCase__					)
							model.eval()
							lowerCAmelCase_      =     model(UpperCamelCase__    ,    attention_mask=UpperCamelCase__    ,    labels=UpperCamelCase__					)
							self.parent.assertEqual(result.logits.shape    ,    (self.batch_size, self.num_labels)					)




				def 					__a						( self    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase					)    ->				Optional[int]:
							lowerCAmelCase_      =     self.num_labels
							lowerCAmelCase_      =     ErnieMForTokenClassification(config=UpperCamelCase__					)
							model.to(UpperCamelCase__					)
							model.eval()
							input_ids.to(UpperCamelCase__					)
							input_mask.to(UpperCamelCase__					)
							token_labels.to(UpperCamelCase__					)

							lowerCAmelCase_      =     model(UpperCamelCase__    ,    attention_mask=UpperCamelCase__    ,    labels=UpperCamelCase__					)

							self.parent.assertEqual(result.logits.shape    ,    (self.batch_size, self.seq_length, self.num_labels)					)




				def 					__a						( self    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase					)    ->				Optional[int]:
							lowerCAmelCase_      =     self.num_choices
							lowerCAmelCase_      =     ErnieMForMultipleChoice(config=UpperCamelCase__					)
							model.to(UpperCamelCase__					)
							model.eval()
							lowerCAmelCase_      =     input_ids.unsqueeze(1					).expand(-1    ,    self.num_choices    ,    -1					).contiguous()
							lowerCAmelCase_      =     input_mask.unsqueeze(1					).expand(-1    ,    self.num_choices    ,    -1					).contiguous()
							lowerCAmelCase_      =     model(
							    UpperCamelCase__    ,    attention_mask=UpperCamelCase__    ,    labels=UpperCamelCase__    ,    )
							self.parent.assertEqual(result.logits.shape    ,    (self.batch_size, self.num_choices)					)




				def 					__a						( self					)    ->				Any:
							lowerCAmelCase_      =     self.prepare_config_and_inputs()
							(
							    lowerCAmelCase_				
							)      =     config_and_inputs
							lowerCAmelCase_      =     {"input_ids": input_ids, "attention_mask": input_mask}
							return config, inputs_dict





@require_torch
class      _lowerCAmelCase (     lowerCAmelCase__ ,				lowerCAmelCase__ ,				unittest.TestCase	):
				_lowercase      		=(
				    (
				        ErnieMModel,
				        ErnieMForMultipleChoice,
				        ErnieMForQuestionAnswering,
				        ErnieMForSequenceClassification,
				        ErnieMForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				_lowercase      		=()
				_lowercase      		=(
				    {
				        """feature-extraction""": ErnieMModel,
				        """question-answering""": ErnieMForQuestionAnswering,
				        """text-classification""": ErnieMForSequenceClassification,
				        """token-classification""": ErnieMForTokenClassification,
				        """zero-shot""": ErnieMForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				_lowercase      		=False




				def 					__a						( self    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase    ,    _UpperCamelCase					)    ->				List[Any]:
							if pipeline_test_casse_name == "QAPipelineTests":
										return True

							return False




				def 					__a						( self					)    ->				Tuple:
							lowerCAmelCase_      =     ErnieMModelTester(self					)
							lowerCAmelCase_      =     ConfigTester(self    ,    config_class=UpperCamelCase__    ,    hidden_size=37					)




				def 					__a						( self					)    ->				List[Any]:
							self.config_tester.run_common_tests()




				def 					__a						( self					)    ->				List[str]:
							lowerCAmelCase_      =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*UpperCamelCase__					)




				def 					__a						( self					)    ->				int:
							lowerCAmelCase_      =     self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
										lowerCAmelCase_      =     type
										self.model_tester.create_and_check_model(*UpperCamelCase__					)




				def 					__a						( self					)    ->				Dict:
							lowerCAmelCase_      =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase__					)




				def 					__a						( self					)    ->				List[Any]:
							lowerCAmelCase_      =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*UpperCamelCase__					)




				def 					__a						( self					)    ->				Any:
							lowerCAmelCase_      =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase__					)




				def 					__a						( self					)    ->				Dict:
							lowerCAmelCase_      =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_information_extraction(*UpperCamelCase__					)




				def 					__a						( self					)    ->				Dict:
							lowerCAmelCase_      =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*UpperCamelCase__					)




				@slow
				def 					__a						( self					)    ->				Optional[Any]:
							for model_name in ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCAmelCase_      =     ErnieMModel.from_pretrained(UpperCamelCase__					)
										self.assertIsNotNone(UpperCamelCase__					)





@require_torch
class      _lowerCAmelCase (     unittest.TestCase	):




				@slow
				def 					__a						( self					)    ->				Dict:
							lowerCAmelCase_      =     ErnieMModel.from_pretrained("susnato/ernie-m-base_pytorch"					)
							model.eval()
							lowerCAmelCase_      =     torch.tensor([[0, 1, 2, 3, 4, 5]]					)
							lowerCAmelCase_      =     model(UpperCamelCase__					)[0]

							# TODO Replace vocab size
							lowerCAmelCase_      =     768

							lowerCAmelCase_      =     torch.Size((1, 6, hidden_size)					)
							self.assertEqual(output.shape    ,    UpperCamelCase__					)

							lowerCAmelCase_      =     torch.tensor(
							    [[[-0.0012, 0.1245, -0.0214], [-0.0742, 0.0244, -0.0771], [-0.0333, 0.1164, -0.1554]]]					)

							self.assertTrue(torch.allclose(output[:, :3, :3]    ,    UpperCamelCase__    ,    atol=1e-3					)					)





import torch
import torch.nn.functional as F

from diffusers import VQDiffusionScheduler

from .test_schedulers import SchedulerCommonTest





class      _lowerCAmelCase (     __a	):
				_lowercase      		=(VQDiffusionScheduler,)




				def 					__a						( self    ,    **_UpperCamelCase					)    ->				List[str]:
							lowerCAmelCase_      =     {
							    "num_vec_classes": 4_097,
							    "num_train_timesteps": 100,
							}

							config.update(**_UpperCamelCase					)
							return config




				def 					__a						( self    ,    _UpperCamelCase					)    ->				str:
							lowerCAmelCase_      =     4
							lowerCAmelCase_      =     8
							lowerCAmelCase_      =     8

							lowerCAmelCase_      =     torch.randint(0    ,    _UpperCamelCase    ,    (batch_size, height * width)					)

							return sample




				@property
				def 					__a						( self					)    ->				List[str]:
							assert False




				def 					__a						( self    ,    _UpperCamelCase					)    ->				Tuple:
							def model(_UpperCamelCase    ,    _UpperCamelCase    ,    *_UpperCamelCase					):
										lowerCAmelCase_     ,  lowerCAmelCase_      =     sample.shape
										lowerCAmelCase_      =     torch.rand((batch_size, num_vec_classes - 1, num_latent_pixels)					)
										lowerCAmelCase_      =     F.log_softmax(logits.double()    ,    dim=1					).float()
										return return_value

							return model




				def 					__a						( self					)    ->				str:
							for timesteps in [2, 5, 100, 1_000]:
										self.check_over_configs(num_train_timesteps=_UpperCamelCase					)




				def 					__a						( self					)    ->				Dict:
							for num_vec_classes in [5, 100, 1_000, 4_000]:
										self.check_over_configs(num_vec_classes=_UpperCamelCase					)




				def 					__a						( self					)    ->				Tuple:
							for t in [0, 50, 99]:
										self.check_over_forward(time_step=_UpperCamelCase					)




				def 					__a						( self					)    ->				List[str]:
							pass



