'''simple docstring'''







from math import acos, sin
from typing import List, Tuple, Union

import numpy as np
import torch
from PIL import Image

from ...models import AutoencoderKL, UNetaDConditionModel
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import AudioPipelineOutput, BaseOutput, DiffusionPipeline, ImagePipelineOutput
from .mel import Mel




class _SCREAMING_SNAKE_CASE(							_SCREAMING_SNAKE_CASE					):
						A_			:							Dict     		= ['vqvae']




						def __init__(    self :       Union[str, Any]							,							UpperCamelCase_ :       AutoencoderKL							,							UpperCamelCase_ :       UNetaDConditionModel							,							UpperCamelCase_ :       Mel							,							UpperCamelCase_ :       Union[DDIMScheduler, DDPMScheduler]							,							)     ->				Optional[Any]:
													super().__init__()
													self.register_modules(unet=UpperCamelCase_							,							scheduler=UpperCamelCase_							,							mel=UpperCamelCase_							,							vqvae=UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       List[Any]      )     ->				int:
													return 50 if isinstance(self.scheduler							,							UpperCamelCase_      ) else 10_00




						@torch.no_grad()
						def __call__(    self :       Dict							,							UpperCamelCase_ :       int = 1							,							UpperCamelCase_ :       str = None							,							UpperCamelCase_ :       np.ndarray = None							,							UpperCamelCase_ :       int = 0							,							UpperCamelCase_ :       int = 0							,							UpperCamelCase_ :       int = None							,							UpperCamelCase_ :       torch.Generator = None							,							UpperCamelCase_ :       float = 0							,							UpperCamelCase_ :       float = 0							,							UpperCamelCase_ :       torch.Generator = None							,							UpperCamelCase_ :       float = 0							,							UpperCamelCase_ :       torch.Tensor = None							,							UpperCamelCase_ :       torch.Tensor = None							,							UpperCamelCase_ :       Dict=True							,							)     ->				Union[
						    Union[AudioPipelineOutput, ImagePipelineOutput],
						    Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]],
						]:
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = steps or self.get_default_steps()
													self.scheduler.set_timesteps(UpperCamelCase_      )
													SCREAMING_SNAKE_CASE__     :Any						    = step_generator or generator
													# For backwards compatibility
													if type(self.unet.config.sample_size      ) == int:
																				SCREAMING_SNAKE_CASE__     :str						    = (self.unet.config.sample_size, self.unet.config.sample_size)
													if noise is None:
																				SCREAMING_SNAKE_CASE__     :int						    = randn_tensor(
																				    (
																				        batch_size,
																				        self.unet.config.in_channels,
																				        self.unet.config.sample_size[0],
																				        self.unet.config.sample_size[1],
																				    )							,							generator=UpperCamelCase_							,							device=self.device							,							)
													SCREAMING_SNAKE_CASE__     :List[Any]						    = noise
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = None

													if audio_file is not None or raw_audio is not None:
																				self.mel.load_audio(UpperCamelCase_							,							UpperCamelCase_      )
																				SCREAMING_SNAKE_CASE__     :List[Any]						    = self.mel.audio_slice_to_image(UpperCamelCase_      )
																				SCREAMING_SNAKE_CASE__     :Any						    = np.frombuffer(input_image.tobytes()							,							dtype='uint8'      ).reshape(
																				    (input_image.height, input_image.width)      )
																				SCREAMING_SNAKE_CASE__     :Any						    = (input_image / 2_55) * 2 - 1
																				SCREAMING_SNAKE_CASE__     :Dict						    = torch.tensor(input_image[np.newaxis, :, :]							,							dtype=torch.float      ).to(self.device      )

																				if self.vqvae is not None:
																											SCREAMING_SNAKE_CASE__     :List[Any]						    = self.vqvae.encode(torch.unsqueeze(UpperCamelCase_							,							0      )      ).latent_dist.sample(
																											    generator=UpperCamelCase_      )[0]
																											SCREAMING_SNAKE_CASE__     :List[Any]						    = self.vqvae.config.scaling_factor * input_images

																				if start_step > 0:
																											SCREAMING_SNAKE_CASE__     :Optional[Any]						    = self.scheduler.add_noise(UpperCamelCase_							,							UpperCamelCase_							,							self.scheduler.timesteps[start_step - 1]      )

																				SCREAMING_SNAKE_CASE__     :Any						    = (
																				    self.unet.config.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length
																				)
																				SCREAMING_SNAKE_CASE__     :Optional[int]						    = int(mask_start_secs * pixels_per_second      )
																				SCREAMING_SNAKE_CASE__     :Tuple						    = int(mask_end_secs * pixels_per_second      )
																				SCREAMING_SNAKE_CASE__     :int						    = self.scheduler.add_noise(UpperCamelCase_							,							UpperCamelCase_							,							torch.tensor(self.scheduler.timesteps[start_step:]      )      )

													for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:]      )      ):
																				if isinstance(self.unet							,							UpperCamelCase_      ):
																											SCREAMING_SNAKE_CASE__     :str						    = self.unet(UpperCamelCase_							,							UpperCamelCase_							,							UpperCamelCase_      )['sample']
																				else:
																											SCREAMING_SNAKE_CASE__     :List[Any]						    = self.unet(UpperCamelCase_							,							UpperCamelCase_      )['sample']

																				if isinstance(self.scheduler							,							UpperCamelCase_      ):
																											SCREAMING_SNAKE_CASE__     :Any						    = self.scheduler.step(
																											    model_output=UpperCamelCase_							,							timestep=UpperCamelCase_							,							sample=UpperCamelCase_							,							eta=UpperCamelCase_							,							generator=UpperCamelCase_							,							)['prev_sample']
																				else:
																											SCREAMING_SNAKE_CASE__     :str						    = self.scheduler.step(
																											    model_output=UpperCamelCase_							,							timestep=UpperCamelCase_							,							sample=UpperCamelCase_							,							generator=UpperCamelCase_							,							)['prev_sample']

																				if mask is not None:
																											if mask_start > 0:
																																		SCREAMING_SNAKE_CASE__     :Any						    = mask[:, step, :, :mask_start]
																											if mask_end > 0:
																																		SCREAMING_SNAKE_CASE__     :Tuple						    = mask[:, step, :, -mask_end:]

													if self.vqvae is not None:
																				# 0.18215 was scaling factor used in training to ensure unit variance
																				SCREAMING_SNAKE_CASE__     :Any						    = 1 / self.vqvae.config.scaling_factor * images
																				SCREAMING_SNAKE_CASE__     :Any						    = self.vqvae.decode(UpperCamelCase_      )['sample']

													SCREAMING_SNAKE_CASE__     :Dict						    = (images / 2 + 0.5).clamp(0							,							1      )
													SCREAMING_SNAKE_CASE__     :str						    = images.cpu().permute(0							,							2							,							3							,							1      ).numpy()
													SCREAMING_SNAKE_CASE__     :List[str]						    = (images * 2_55).round().astype('uint8'      )
													SCREAMING_SNAKE_CASE__     :Any						    = list(
													    (Image.fromarray(_[:, :, 0]      ) for _ in images)
													    if images.shape[3] == 1
													    else (Image.fromarray(UpperCamelCase_							,							mode='RGB'      ).convert('L'      ) for _ in images)      )

													SCREAMING_SNAKE_CASE__     :Optional[int]						    = [self.mel.image_to_audio(UpperCamelCase_      ) for _ in images]
													if not return_dict:
																				return images, (self.mel.get_sample_rate(), audios)

													return BaseOutput(**AudioPipelineOutput(np.array(UpperCamelCase_      )[:, np.newaxis, :]      )							,							**ImagePipelineOutput(UpperCamelCase_      )      )




						@torch.no_grad()
						def      __lowerCamelCase       (    self :       str							,							UpperCamelCase_ :       List[Image.Image]							,							UpperCamelCase_ :       int = 50      )     ->				np.ndarray:
													assert isinstance(self.scheduler							,							UpperCamelCase_      )
													self.scheduler.set_timesteps(UpperCamelCase_      )
													SCREAMING_SNAKE_CASE__     :List[Any]						    = np.array(
													    [np.frombuffer(image.tobytes()							,							dtype='uint8'      ).reshape((1, image.height, image.width)      ) for image in images]      )
													SCREAMING_SNAKE_CASE__     :Tuple						    = (sample / 2_55) * 2 - 1
													SCREAMING_SNAKE_CASE__     :List[str]						    = torch.Tensor(UpperCamelCase_      ).to(self.device      )

													for t in self.progress_bar(torch.flip(self.scheduler.timesteps							,							(0,)      )      ):
																				SCREAMING_SNAKE_CASE__     :Optional[Any]						    = t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
																				SCREAMING_SNAKE_CASE__     :List[str]						    = self.scheduler.alphas_cumprod[t]
																				SCREAMING_SNAKE_CASE__     :Tuple						    = (
																				    self.scheduler.alphas_cumprod[prev_timestep]
																				    if prev_timestep >= 0
																				    else self.scheduler.final_alpha_cumprod
																				)
																				SCREAMING_SNAKE_CASE__     :Optional[Any]						    = 1 - alpha_prod_t
																				SCREAMING_SNAKE_CASE__     :Any						    = self.unet(UpperCamelCase_							,							UpperCamelCase_      )['sample']
																				SCREAMING_SNAKE_CASE__     :Optional[Any]						    = (1 - alpha_prod_t_prev) ** 0.5 * model_output
																				SCREAMING_SNAKE_CASE__     :Optional[int]						    = (sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5)
																				SCREAMING_SNAKE_CASE__     :Tuple						    = sample * alpha_prod_t ** 0.5 + beta_prod_t ** 0.5 * model_output

													return sample




						@staticmethod
						def      __lowerCamelCase       (    UpperCamelCase_ :       torch.Tensor							,							UpperCamelCase_ :       torch.Tensor							,							UpperCamelCase_ :       float      )     ->				torch.Tensor:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = acos(torch.dot(torch.flatten(UpperCamelCase_      )							,							torch.flatten(UpperCamelCase_      )      ) / torch.norm(UpperCamelCase_      ) / torch.norm(UpperCamelCase_      )      )
													return sin((1 - alpha) * theta      ) * xa / sin(UpperCamelCase_      ) + sin(alpha * theta      ) * xa / sin(UpperCamelCase_      )







'''simple docstring'''







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




						def __init__(    self :       Optional[int]							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       Tuple=13							,							UpperCamelCase_ :       Union[str, Any]=7							,							UpperCamelCase_ :       Union[str, Any]=True							,							UpperCamelCase_ :       Union[str, Any]=True							,							UpperCamelCase_ :       Any=True							,							UpperCamelCase_ :       int=True							,							UpperCamelCase_ :       List[Any]=99							,							UpperCamelCase_ :       List[str]=64							,							UpperCamelCase_ :       List[str]=32							,							UpperCamelCase_ :       Union[str, Any]=5							,							UpperCamelCase_ :       Dict=4							,							UpperCamelCase_ :       List[str]=37							,							UpperCamelCase_ :       str="gelu"							,							UpperCamelCase_ :       List[Any]=0.1							,							UpperCamelCase_ :       Union[str, Any]=0.1							,							UpperCamelCase_ :       Optional[Any]=5_12							,							UpperCamelCase_ :       Union[str, Any]=16							,							UpperCamelCase_ :       Union[str, Any]=2							,							UpperCamelCase_ :       Optional[int]=0.02							,							UpperCamelCase_ :       Optional[Any]=3							,							UpperCamelCase_ :       Any=4							,							UpperCamelCase_ :       int=None							,							)     ->				Optional[int]:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = parent
													SCREAMING_SNAKE_CASE__     :Any						    = batch_size
													SCREAMING_SNAKE_CASE__     :Tuple						    = seq_length
													SCREAMING_SNAKE_CASE__     :Any						    = is_training
													SCREAMING_SNAKE_CASE__     :int						    = use_input_mask
													SCREAMING_SNAKE_CASE__     :Dict						    = use_token_type_ids
													SCREAMING_SNAKE_CASE__     :str						    = use_labels
													SCREAMING_SNAKE_CASE__     :str						    = vocab_size
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = hidden_size
													SCREAMING_SNAKE_CASE__     :List[Any]						    = embedding_size
													SCREAMING_SNAKE_CASE__     :int						    = num_hidden_layers
													SCREAMING_SNAKE_CASE__     :int						    = num_attention_heads
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = intermediate_size
													SCREAMING_SNAKE_CASE__     :str						    = hidden_act
													SCREAMING_SNAKE_CASE__     :Any						    = hidden_dropout_prob
													SCREAMING_SNAKE_CASE__     :Any						    = attention_probs_dropout_prob
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = max_position_embeddings
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = type_vocab_size
													SCREAMING_SNAKE_CASE__     :Tuple						    = type_sequence_label_size
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = initializer_range
													SCREAMING_SNAKE_CASE__     :int						    = num_labels
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = num_choices
													SCREAMING_SNAKE_CASE__     :List[Any]						    = scope




						def      __lowerCamelCase       (    self :       Any      )     ->				Optional[Any]:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = ids_tensor([self.batch_size, self.seq_length]							,							self.vocab_size      )

													SCREAMING_SNAKE_CASE__     :List[Any]						    = None
													if self.use_input_mask:
																				SCREAMING_SNAKE_CASE__     :str						    = random_attention_mask([self.batch_size, self.seq_length]      )

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

													SCREAMING_SNAKE_CASE__     :List[str]						    = None
													SCREAMING_SNAKE_CASE__     :int						    = None
													SCREAMING_SNAKE_CASE__     :Tuple						    = None
													if self.use_labels:
																				SCREAMING_SNAKE_CASE__     :Any						    = ids_tensor([self.batch_size]							,							self.type_sequence_label_size      )
																				SCREAMING_SNAKE_CASE__     :Optional[Any]						    = ids_tensor([self.batch_size, self.seq_length]							,							self.num_labels      )
																				SCREAMING_SNAKE_CASE__     :Any						    = ids_tensor([self.batch_size]							,							self.num_choices      )

													SCREAMING_SNAKE_CASE__     :str						    = self.get_config()

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




						def      __lowerCamelCase       (    self :       List[str]      )     ->				Any:
													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      __lowerCamelCase       (    self :       Optional[int]							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       str							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       int      )     ->				int:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = MegatronBertModel(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :int						    = model(UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_      )
													SCREAMING_SNAKE_CASE__     :str						    = model(UpperCamelCase_							,							token_type_ids=UpperCamelCase_      )
													SCREAMING_SNAKE_CASE__     :List[str]						    = 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      __lowerCamelCase       (    self :       str							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       Optional[Any]							,							UpperCamelCase_ :       str							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       Union[str, Any]      )     ->				List[str]:
													SCREAMING_SNAKE_CASE__     :str						    = MegatronBertForMaskedLM(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = 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      __lowerCamelCase       (    self :       Any							,							UpperCamelCase_ :       List[Any]							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       int      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = MegatronBertForCausalLM(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :List[str]						    = 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      __lowerCamelCase       (    self :       List[str]							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       Optional[Any]							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       Union[str, Any]							,							UpperCamelCase_ :       Optional[int]      )     ->				List[Any]:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = MegatronBertForNextSentencePrediction(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :List[str]						    = model(
													    UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							labels=UpperCamelCase_							,							)
													self.parent.assertEqual(result.logits.shape							,							(self.batch_size, 2)      )




						def      __lowerCamelCase       (    self :       Dict							,							UpperCamelCase_ :       List[Any]							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       str							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       Tuple							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       Optional[int]      )     ->				Any:
													SCREAMING_SNAKE_CASE__     :Tuple						    = MegatronBertForPreTraining(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = 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      __lowerCamelCase       (    self :       Union[str, Any]							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       Tuple							,							UpperCamelCase_ :       List[Any]							,							UpperCamelCase_ :       List[Any]      )     ->				Tuple:
													SCREAMING_SNAKE_CASE__     :int						    = MegatronBertForQuestionAnswering(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :Tuple						    = 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      __lowerCamelCase       (    self :       List[Any]							,							UpperCamelCase_ :       Tuple							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       str							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       Tuple							,							UpperCamelCase_ :       Any      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :List[Any]						    = self.num_labels
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = MegatronBertForSequenceClassification(UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :int						    = model(UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							labels=UpperCamelCase_      )
													self.parent.assertEqual(result.logits.shape							,							(self.batch_size, self.num_labels)      )




						def      __lowerCamelCase       (    self :       Any							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       Union[str, Any]							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       Union[str, Any]							,							UpperCamelCase_ :       Union[str, Any]							,							UpperCamelCase_ :       Optional[Any]      )     ->				Any:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.num_labels
													SCREAMING_SNAKE_CASE__     :Any						    = MegatronBertForTokenClassification(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = 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      __lowerCamelCase       (    self :       str							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       List[Any]							,							UpperCamelCase_ :       Union[str, Any]							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       str							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       List[Any]      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = self.num_choices
													SCREAMING_SNAKE_CASE__     :Tuple						    = MegatronBertForMultipleChoice(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :Any						    = input_ids.unsqueeze(1      ).expand(-1							,							self.num_choices							,							-1      ).contiguous()
													SCREAMING_SNAKE_CASE__     :List[Any]						    = token_type_ids.unsqueeze(1      ).expand(-1							,							self.num_choices							,							-1      ).contiguous()
													SCREAMING_SNAKE_CASE__     :Dict						    = input_mask.unsqueeze(1      ).expand(-1							,							self.num_choices							,							-1      ).contiguous()
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = model(
													    UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							labels=UpperCamelCase_							,							)
													self.parent.assertEqual(result.logits.shape							,							(self.batch_size, self.num_choices)      )




						def      __lowerCamelCase       (    self :       Any      )     ->				int:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.prepare_config_and_inputs()
													(
													    (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      
													)     :Union[str, Any]						    = config_and_inputs
													SCREAMING_SNAKE_CASE__     :Tuple						    = {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
													return config, inputs_dict




@require_torch
class _SCREAMING_SNAKE_CASE(							_SCREAMING_SNAKE_CASE						,						_SCREAMING_SNAKE_CASE						,						unittest.TestCase					):
						A_			:							List[Any]     		= (
						    (
						        MegatronBertModel,
						        MegatronBertForMaskedLM,
						        MegatronBertForCausalLM,
						        MegatronBertForMultipleChoice,
						        MegatronBertForNextSentencePrediction,
						        MegatronBertForPreTraining,
						        MegatronBertForQuestionAnswering,
						        MegatronBertForSequenceClassification,
						        MegatronBertForTokenClassification,
						    )
						    if is_torch_available()
						    else ()
						)
						A_			:							List[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 {}
						)
						A_			:							int     		= True
						# test_resize_embeddings = False
						A_			:							Dict     		= False




						def      __lowerCamelCase       (    self :       Union[str, Any]							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       str=False      )     ->				List[Any]:
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = super()._prepare_for_class(UpperCamelCase_							,							UpperCamelCase_							,							return_labels=UpperCamelCase_      )

													if return_labels:
																				if model_class in get_values(UpperCamelCase_      ):
																											SCREAMING_SNAKE_CASE__     :str						    = torch.zeros(
																											    (self.model_tester.batch_size, self.model_tester.seq_length)							,							dtype=torch.long							,							device=UpperCamelCase_      )
																											SCREAMING_SNAKE_CASE__     :List[str]						    = torch.zeros(
																											    self.model_tester.batch_size							,							dtype=torch.long							,							device=UpperCamelCase_      )
													return inputs_dict




						def      __lowerCamelCase       (    self :       Union[str, Any]      )     ->				Any:
													SCREAMING_SNAKE_CASE__     :Dict						    = MegatronBertModelTester(self      )
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = ConfigTester(self							,							config_class=UpperCamelCase_							,							hidden_size=37      )




						def      __lowerCamelCase       (    self :       List[str]      )     ->				Optional[int]:
													self.config_tester.run_common_tests()




						def      __lowerCamelCase       (    self :       List[str]      )     ->				List[str]:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_model(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       Any      )     ->				Tuple:
													SCREAMING_SNAKE_CASE__     :List[Any]						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_masked_lm(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       Union[str, Any]      )     ->				str:
													SCREAMING_SNAKE_CASE__     :Any						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       str      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       int      )     ->				Any:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_pretraining(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       Dict      )     ->				Union[str, Any]:
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_question_answering(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       Any      )     ->				Union[str, Any]:
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       List[Any]      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_token_classification(*UpperCamelCase_      )


def    lowerCamelCase       (			UpperCAmelCase__				:  Dict		)			->	Any:


							'''simple docstring'''






							return torch.tensor(
							    UpperCAmelCase__				,						dtype=torch.long				,						device=UpperCAmelCase__				,						)


UpperCamelCase_      	=	1e-4





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




						@slow
						@unittest.skip('Model is not available.'      )
						def      __lowerCamelCase       (    self :       str      )     ->				int:
													SCREAMING_SNAKE_CASE__     :Tuple						    = 'nvidia/megatron-bert-uncased-345m'
													if "MYDIR" in os.environ:
																				SCREAMING_SNAKE_CASE__     :Dict						    = os.path.join(os.environ['MYDIR']							,							UpperCamelCase_      )
													SCREAMING_SNAKE_CASE__     :Dict						    = MegatronBertModel.from_pretrained(UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.half()
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = _long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]      )
													with torch.no_grad():
																				SCREAMING_SNAKE_CASE__     :str						    = model(UpperCamelCase_      )[0]
													SCREAMING_SNAKE_CASE__     :Dict						    = torch.Size((1, 9, 10_24)      )
													self.assertEqual(output.shape							,							UpperCamelCase_      )

													SCREAMING_SNAKE_CASE__     :Tuple						    = [-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
													for ii in range(3      ):
																				for jj in range(3      ):
																											SCREAMING_SNAKE_CASE__     :List[Any]						    = output[0, ii, jj]
																											SCREAMING_SNAKE_CASE__     :List[str]						    = expected[3 * ii + jj]
																											SCREAMING_SNAKE_CASE__     :List[Any]						    = 'ii={} jj={} a={} b={}'.format(UpperCamelCase_							,							UpperCamelCase_							,							UpperCamelCase_							,							UpperCamelCase_      )
																											self.assertTrue(math.isclose(UpperCamelCase_							,							UpperCamelCase_							,							rel_tol=UpperCamelCase_							,							abs_tol=UpperCamelCase_      )							,							msg=UpperCamelCase_      )






