



import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, 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,
						    MegatronBertForCausalLM,
						    MegatronBertForMaskedLM,
						    MegatronBertForMultipleChoice,
						    MegatronBertForNextSentencePrediction,
						    MegatronBertForPreTraining,
						    MegatronBertForQuestionAnswering,
						    MegatronBertForSequenceClassification,
						    MegatronBertForTokenClassification,
						    MegatronBertModel,
						)

class  __snake_case      :







							'''simple docstring'''






							def __init__(							self  ,  a_  ,  a_=13  ,  a_=7  ,  a_=True  ,  a_=True  ,  a_=True  ,  a_=True  ,  a_=99  ,  a_=64  ,  a_=32  ,  a_=5  ,  a_=4  ,  a_=37  ,  a_="gelu"  ,  a_=0.1  ,  a_=0.1  ,  a_=512  ,  a_=16  ,  a_=2  ,  a_=0.02  ,  a_=3  ,  a_=4  ,  a_=None  ,  ):
											a__				      =		parent
											a__				      =		batch_size
											a__				      =		seq_length
											a__				      =		is_training
											a__				      =		use_input_mask
											a__				      =		use_token_type_ids
											a__				      =		use_labels
											a__				      =		vocab_size
											a__				      =		hidden_size
											a__				      =		embedding_size
											a__				      =		num_hidden_layers
											a__				      =		num_attention_heads
											a__				      =		intermediate_size
											a__				      =		hidden_act
											a__				      =		hidden_dropout_prob
											a__				      =		attention_probs_dropout_prob
											a__				      =		max_position_embeddings
											a__				      =		type_vocab_size
											a__				      =		type_sequence_label_size
											a__				      =		initializer_range
											a__				      =		num_labels
											a__				      =		num_choices
											a__				      =		scope


							def _a      (							self		):
											a__				      =		ids_tensor([self.batch_size, self.seq_length]  ,  self.vocab_size		)

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

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

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

											a__				      =		self.get_config()

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


							def _a      (							self		):
											return MegatronBertConfig(
											    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  ,  embedding_size=self.embedding_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  ,  is_decoder=__UpperCamelCase  ,  initializer_range=self.initializer_range  ,  )


							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		MegatronBertModel(config=__UpperCamelCase		)
											model.to(__UpperCamelCase		)
											model.eval()
											a__				      =		model(__UpperCamelCase  ,  attention_mask=__UpperCamelCase  ,  token_type_ids=__UpperCamelCase		)
											a__				      =		model(__UpperCamelCase  ,  token_type_ids=__UpperCamelCase		)
											a__				      =		model(__UpperCamelCase		)

											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_		):
											a__				      =		MegatronBertForMaskedLM(config=__UpperCamelCase		)
											model.to(__UpperCamelCase		)
											model.eval()
											a__				      =		model(__UpperCamelCase  ,  attention_mask=__UpperCamelCase  ,  token_type_ids=__UpperCamelCase  ,  labels=__UpperCamelCase		)
											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_		):
											a__				      =		MegatronBertForCausalLM(config=__UpperCamelCase		)
											model.to(__UpperCamelCase		)
											model.eval()
											a__				      =		model(__UpperCamelCase  ,  attention_mask=__UpperCamelCase  ,  token_type_ids=__UpperCamelCase  ,  labels=__UpperCamelCase		)
											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_		):
											a__				      =		MegatronBertForNextSentencePrediction(config=__UpperCamelCase		)
											model.to(__UpperCamelCase		)
											model.eval()
											a__				      =		model(
											    __UpperCamelCase  ,  attention_mask=__UpperCamelCase  ,  token_type_ids=__UpperCamelCase  ,  labels=__UpperCamelCase  ,  )
											self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, 2)		)


							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		MegatronBertForPreTraining(config=__UpperCamelCase		)
											model.to(__UpperCamelCase		)
											model.eval()
											a__				      =		model(
											    __UpperCamelCase  ,  attention_mask=__UpperCamelCase  ,  token_type_ids=__UpperCamelCase  ,  labels=__UpperCamelCase  ,  next_sentence_label=__UpperCamelCase  ,  )
											self.parent.assertEqual(result.prediction_logits.shape  ,  (self.batch_size, self.seq_length, self.vocab_size)		)
											self.parent.assertEqual(result.seq_relationship_logits.shape  ,  (self.batch_size, 2)		)


							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		MegatronBertForQuestionAnswering(config=__UpperCamelCase		)
											model.to(__UpperCamelCase		)
											model.eval()
											a__				      =		model(
											    __UpperCamelCase  ,  attention_mask=__UpperCamelCase  ,  token_type_ids=__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  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		self.num_labels
											a__				      =		MegatronBertForSequenceClassification(__UpperCamelCase		)
											model.to(__UpperCamelCase		)
											model.eval()
											a__				      =		model(__UpperCamelCase  ,  attention_mask=__UpperCamelCase  ,  token_type_ids=__UpperCamelCase  ,  labels=__UpperCamelCase		)
											self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_labels)		)


							def _a      (							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											a__				      =		self.num_labels
											a__				      =		MegatronBertForTokenClassification(config=__UpperCamelCase		)
											model.to(__UpperCamelCase		)
											model.eval()
											a__				      =		model(__UpperCamelCase  ,  attention_mask=__UpperCamelCase  ,  token_type_ids=__UpperCamelCase  ,  labels=__UpperCamelCase		)
											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_		):
											a__				      =		self.num_choices
											a__				      =		MegatronBertForMultipleChoice(config=__UpperCamelCase		)
											model.to(__UpperCamelCase		)
											model.eval()
											a__				      =		input_ids.unsqueeze(1		).expand(-1  ,  self.num_choices  ,  -1		).contiguous()
											a__				      =		token_type_ids.unsqueeze(1		).expand(-1  ,  self.num_choices  ,  -1		).contiguous()
											a__				      =		input_mask.unsqueeze(1		).expand(-1  ,  self.num_choices  ,  -1		).contiguous()
											a__				      =		model(
											    __UpperCamelCase  ,  attention_mask=__UpperCamelCase  ,  token_type_ids=__UpperCamelCase  ,  labels=__UpperCamelCase  ,  )
											self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_choices)		)







							def _a      (							self		):
											a__				      =		self.prepare_config_and_inputs()
											(
											    (
											    a__       
											)      ,			(
											    a__       
											)      ,			(
											    a__       
											)      ,			(
											    a__       
											)      ,			(
											    a__       
											)      ,			(
											    a__       
											)      ,			(
											    a__       
											)      ,			       
											)				      =		config_and_inputs
											a__				      =		{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
											return config, inputs_dict




@require_torch
class  __snake_case      (   __snake_case			,__snake_case			,unittest.TestCase):







							'''simple docstring'''






							UpperCamelCase__				:      str         =					(
							    (
							        MegatronBertModel,
							        MegatronBertForMaskedLM,
							        MegatronBertForCausalLM,
							        MegatronBertForMultipleChoice,
							        MegatronBertForNextSentencePrediction,
							        MegatronBertForPreTraining,
							        MegatronBertForQuestionAnswering,
							        MegatronBertForSequenceClassification,
							        MegatronBertForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							UpperCamelCase__				:      Any         =					(
							    {
							        """feature-extraction""": MegatronBertModel,
							        """fill-mask""": MegatronBertForMaskedLM,
							        """question-answering""": MegatronBertForQuestionAnswering,
							        """text-classification""": MegatronBertForSequenceClassification,
							        """text-generation""": MegatronBertForCausalLM,
							        """token-classification""": MegatronBertForTokenClassification,
							        """zero-shot""": MegatronBertForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							UpperCamelCase__				:      Any         =					True
							# test_resize_embeddings = False
							UpperCamelCase__				:      int         =					False


							def _a      (							self  ,  a_  ,  a_  ,  a_=False		):
											a__				      =		super()._prepare_for_class(__UpperCamelCase  ,  __UpperCamelCase  ,  return_labels=__UpperCamelCase		)

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


							def _a      (							self		):
											a__				      =		MegatronBertModelTester(self		)
											a__				      =		ConfigTester(self  ,  config_class=__UpperCamelCase  ,  hidden_size=37		)


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


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_model(*__UpperCamelCase		)


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__UpperCamelCase		)


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__UpperCamelCase		)


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__UpperCamelCase		)


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_pretraining(*__UpperCamelCase		)


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_question_answering(*__UpperCamelCase		)


							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__UpperCamelCase		)







							def _a      (							self		):
											a__				      =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_token_classification(*__UpperCamelCase		)




def 			A_      (	__a     :				str		):

				"""simple docstring"""
				return torch.tensor(
				    lowercase__ ,		dtype=torch.long ,		device=lowercase__ ,		)


UpperCAmelCase							 =					1E-4




@require_torch
@require_sentencepiece
@require_tokenizers
class  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							@slow
							@unittest.skip("""Model is not available."""		)
							def _a      (							self		):
											a__				      =		"""nvidia/megatron-bert-uncased-345m"""
											if "MYDIR" in os.environ:
															a__				      =		os.path.join(os.environ["""MYDIR"""]  ,  __UpperCamelCase		)
											a__				      =		MegatronBertModel.from_pretrained(__UpperCamelCase		)
											model.to(__UpperCamelCase		)
											model.half()
											a__				      =		_long_tensor([[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]]		)
											with torch.no_grad():
															a__				      =		model(__UpperCamelCase		)[0]
											a__				      =		torch.Size((1, 9, 1_024)		)
											self.assertEqual(output.shape  ,  __UpperCamelCase		)

											a__				      =		[-0.6_040, -0.2_517, -0.1_025, 0.3_420, -0.6_758, -0.0_017, -0.1_089, -0.1_990, 0.5_728]
											for ii in range(3		):
															for jj in range(3		):
																			a__				      =		output[0, ii, jj]
																			a__				      =		expected[3 * ii + jj]
																			a__				      =		"""ii={} jj={} a={} b={}""".format(__UpperCamelCase  ,  __UpperCamelCase  ,  __UpperCamelCase  ,  __UpperCamelCase		)
																			self.assertTrue(math.isclose(__UpperCamelCase  ,  __UpperCamelCase  ,  rel_tol=__UpperCamelCase  ,  abs_tol=__UpperCamelCase		)  ,  msg=__UpperCamelCase		)






import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()

class  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							def _a      (							self		):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()


							@property
							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		UNetaDModel(
											    sample_size=(32, 64)  ,  in_channels=1  ,  out_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")  ,  up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")  ,  )
											return model


							@property
							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		UNetaDConditionModel(
											    sample_size=(64, 32)  ,  in_channels=1  ,  out_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")  ,  up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")  ,  cross_attention_dim=10  ,  )
											return model


							@property
							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		AutoencoderKL(
											    sample_size=(128, 64)  ,  in_channels=1  ,  out_channels=1  ,  latent_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""DownEncoderBlock2D""", """DownEncoderBlock2D""")  ,  up_block_types=("""UpDecoderBlock2D""", """UpDecoderBlock2D""")  ,  )
											a__				      =		UNetaDModel(
											    sample_size=(64, 32)  ,  in_channels=1  ,  out_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")  ,  up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")  ,  )
											return vqvae, unet







							@slow
							def _a      (							self		):
											a__				      =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
											a__				      =		Mel(
											    x_res=self.dummy_unet.config.sample_size[1]  ,  y_res=self.dummy_unet.config.sample_size[0]  ,  )

											a__				      =		DDPMScheduler()
											a__				      =		AudioDiffusionPipeline(vqvae=a_  ,  unet=self.dummy_unet  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_  ,  steps=4		)
											a__				      =		output.audios[0]
											a__				      =		output.images[0]

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_  ,  steps=4  ,  return_dict=a_		)
											a__				      =		output[0][0]

											assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
											assert (
											    image.height == self.dummy_unet.config.sample_size[0]
											    and image.width == self.dummy_unet.config.sample_size[1]
											)
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.frombuffer(image_from_tuple.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([69, 255, 255, 255, 0, 0, 77, 181, 12, 127]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0
											assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() == 0

											a__				      =		Mel(
											    x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]  ,  y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]  ,  )

											a__				      =		DDIMScheduler()
											a__				      =		self.dummy_vqvae_and_unet
											a__				      =		AudioDiffusionPipeline(
											    vqvae=self.dummy_vqvae_and_unet[0]  ,  unet=dummy_vqvae_and_unet[1]  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											np.random.seed(0		)
											a__				      =		np.random.uniform(-1  ,  1  ,  ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)		)
											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(raw_audio=a_  ,  generator=a_  ,  start_step=5  ,  steps=10		)
											a__				      =		output.images[0]

											assert (
											    image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
											    and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
											)
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0

											a__				      =		self.dummy_unet_condition
											a__				      =		AudioDiffusionPipeline(
											    vqvae=self.dummy_vqvae_and_unet[0]  ,  unet=a_  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											np.random.seed(0		)
											a__				      =		torch.rand((1, 1, 10)		)
											a__				      =		pipe(generator=a_  ,  encoding=a_		)
											a__				      =		output.images[0]
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0




@slow
@require_torch_gpu
class  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






							def _a      (							self		):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()







							def _a      (							self		):
											a__				      =		torch_device

											a__				      =		DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_		)
											a__				      =		output.audios[0]
											a__				      =		output.images[0]

											assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
											assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]		)

											assert np.abs(image_slice.flatten() - expected_slice		).max() == 0

