





'''simple docstring'''






import math
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, randn_tensor
from .scheduling_utils import SchedulerMixin







@dataclass
class       lowerCamelCase       (    __UpperCAmelCase	):
					_SCREAMING_SNAKE_CASE				 =     42
					_SCREAMING_SNAKE_CASE				 =     42





def    lowercase_   (		lowercase__					,	lowercase__=0.9_9_9					,	lowercase__="cosine"					,	) ->Optional[int]:
						if alpha_transform_type == "cosine":

												def alpha_bar_fn(lowercase__    ):
																		return math.cos((t + 0.0_0_8) / 1.0_0_8 * math.pi / 2    ) ** 2

						elif alpha_transform_type == "exp":

												def alpha_bar_fn(lowercase__    ):
																		return math.exp(t * -1_2.0    )

						else:
												raise ValueError(F'''Unsupported alpha_tranform_type: {alpha_transform_type}'''    )

						_snake_case: str	       =		[]
						for i in range(lowercase__    ):
												_snake_case: Tuple	       =		i / num_diffusion_timesteps
												_snake_case: str	       =		(i + 1) / num_diffusion_timesteps
												betas.append(min(1 - alpha_bar_fn(lowercase__    ) / alpha_bar_fn(lowercase__    )					,	lowercase__    )    )
						return torch.tensor(lowercase__					,	dtype=torch.floataa    )







class       lowerCamelCase       (    __UpperCAmelCase  ,				__UpperCAmelCase	):
					_SCREAMING_SNAKE_CASE				 =     1


					@register_to_config
					def __init__(       self						:	List[Any]							,  __snake_case						:	int = 10_00							,  __snake_case						:	float = 0.0_001							,  __snake_case						:	float = 0.02							,  __snake_case						:	str = "linear"							,  __snake_case						:	float = 0.0							,  __snake_case						:	Optional[np.ndarray] = None							,  __snake_case						:	bool = True							,  ):





											'''simple docstring'''





											if trained_betas is not None:
																	_snake_case: int	       =		torch.from_numpy(__snake_case							)
											elif beta_schedule == "linear":
																	_snake_case: Tuple	       =		torch.linspace(__snake_case							,  __snake_case							,  __snake_case							,  dtype=torch.floataa							)
											elif beta_schedule == "scaled_linear":
																	# this schedule is very specific to the latent diffusion model.
																	_snake_case: Dict	       =		(
																	    torch.linspace(beta_start**0.5							,  beta_end**0.5							,  __snake_case							,  dtype=torch.floataa							) ** 2
																	)
											elif beta_schedule == "squaredcos_cap_v2":
																	# Glide cosine schedule
																	_snake_case: Optional[Any]	       =		betas_for_alpha_bar(__snake_case							)
											elif beta_schedule == "sigmoid":
																	# GeoDiff sigmoid schedule
																	_snake_case: str	       =		torch.linspace(-6							,  6							,  __snake_case							)
																	_snake_case: Optional[int]	       =		torch.sigmoid(__snake_case							) * (beta_end - beta_start) + beta_start
											else:
																	raise NotImplementedError(f'''{beta_schedule} does is not implemented for {self.__class__}'''							)

											_snake_case: int	       =		1.0 - self.betas
											_snake_case: Dict	       =		torch.cumprod(self.alphas							,  dim=0							)
											_snake_case: List[str]	       =		torch.tensor(1.0							)

											_snake_case: Any	       =		torch.tensor(1.0							)

											# standard deviation of the initial noise distribution
											_snake_case: List[Any]	       =		1.0

											# setable values
											_snake_case: Optional[int]	       =		None
											_snake_case: Dict	       =		torch.from_numpy(np.arange(0							,  __snake_case							)[::-1].copy()							)

											_snake_case: Any	       =		eta


					def  SCREAMING_SNAKE_CASE_		(       self						:	int							,  __snake_case						:	torch.FloatTensor							,  __snake_case						:	Optional[int] = None							):





											'''simple docstring'''





											return sample


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[Any]							,  __snake_case						:	int							,  __snake_case						:	int = 10							,  __snake_case						:	int = 10							,  __snake_case						:	Union[str, torch.device] = None							,  ):





											'''simple docstring'''





											_snake_case: Optional[Any]	       =		min(self.config.num_train_timesteps							,  __snake_case							)
											_snake_case: Tuple	       =		num_inference_steps

											_snake_case: List[str]	       =		[]

											_snake_case: Optional[int]	       =		{}
											for j in range(0							,  num_inference_steps - jump_length							,  __snake_case							):
																	_snake_case: Dict	       =		jump_n_sample - 1

											_snake_case: Optional[int]	       =		num_inference_steps
											while t >= 1:
																	_snake_case: Union[str, Any]	       =		t - 1
																	timesteps.append(__snake_case							)

																	if jumps.get(__snake_case							,  0							) > 0:
																							_snake_case: List[Any]	       =		jumps[t] - 1
																							for _ in range(__snake_case							):
																													_snake_case: Optional[int]	       =		t + 1
																													timesteps.append(__snake_case							)

											_snake_case: Any	       =		np.array(__snake_case							) * (self.config.num_train_timesteps // self.num_inference_steps)
											_snake_case: Union[str, Any]	       =		torch.from_numpy(__snake_case							).to(__snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	int							,  __snake_case						:	Any							):





											'''simple docstring'''





											_snake_case: List[str]	       =		t - self.config.num_train_timesteps // self.num_inference_steps

											_snake_case: str	       =		self.alphas_cumprod[t]
											_snake_case: int	       =		self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
											_snake_case: Dict	       =		1 - alpha_prod_t
											_snake_case: List[Any]	       =		1 - alpha_prod_t_prev

											# For t > 0, compute predicted variance βt (see formula (6) and (7) from
											# https://arxiv.org/pdf/2006.11239.pdf) and sample from it to get
											# previous sample x_{t-1} ~ N(pred_prev_sample, variance) == add
											# variance to pred_sample
											# Is equivalent to formula (16) in https://arxiv.org/pdf/2010.02502.pdf
											# without eta.
											# variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * self.betas[t]
											_snake_case: str	       =		(beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)

											return variance


					def  SCREAMING_SNAKE_CASE_		(       self						:	Dict							,  __snake_case						:	torch.FloatTensor							,  __snake_case						:	int							,  __snake_case						:	torch.FloatTensor							,  __snake_case						:	torch.FloatTensor							,  __snake_case						:	torch.FloatTensor							,  __snake_case						:	Optional[torch.Generator] = None							,  __snake_case						:	bool = True							,  ):





											'''simple docstring'''





											_snake_case: Union[str, Any]	       =		timestep
											_snake_case: Dict	       =		timestep - self.config.num_train_timesteps // self.num_inference_steps

											# 1. compute alphas, betas
											_snake_case: Optional[int]	       =		self.alphas_cumprod[t]
											_snake_case: Dict	       =		self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
											_snake_case: Optional[int]	       =		1 - alpha_prod_t

											# 2. compute predicted original sample from predicted noise also called
											# "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
											_snake_case: Optional[Any]	       =		(sample - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.5

											# 3. Clip "predicted x_0"
											if self.config.clip_sample:
																	_snake_case: Optional[int]	       =		torch.clamp(__snake_case							,  -1							,  1							)

											# We choose to follow RePaint Algorithm 1 to get x_{t-1}, however we
											# substitute formula (7) in the algorithm coming from DDPM paper
											# (formula (4) Algorithm 2 - Sampling) with formula (12) from DDIM paper.
											# DDIM schedule gives the same results as DDPM with eta = 1.0
											# Noise is being reused in 7. and 8., but no impact on quality has
											# been observed.

											# 5. Add noise
											_snake_case: Dict	       =		model_output.device
											_snake_case: List[Any]	       =		randn_tensor(model_output.shape							,  generator=__snake_case							,  device=__snake_case							,  dtype=model_output.dtype							)
											_snake_case: Tuple	       =		self.eta * self._get_variance(__snake_case							) ** 0.5

											_snake_case: Union[str, Any]	       =		0
											if t > 0 and self.eta > 0:
																	_snake_case: List[Any]	       =		std_dev_t * noise

											# 6. compute "direction pointing to x_t" of formula (12)
											# from https://arxiv.org/pdf/2010.02502.pdf
											_snake_case: Union[str, Any]	       =		(1 - alpha_prod_t_prev - std_dev_t**2) ** 0.5 * model_output

											# 7. compute x_{t-1} of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
											_snake_case: Dict	       =		alpha_prod_t_prev**0.5 * pred_original_sample + pred_sample_direction + variance

											# 8. Algorithm 1 Line 5 https://arxiv.org/pdf/2201.09865.pdf
											_snake_case: Any	       =		(alpha_prod_t_prev**0.5) * original_image + ((1 - alpha_prod_t_prev) ** 0.5) * noise

											# 9. Algorithm 1 Line 8 https://arxiv.org/pdf/2201.09865.pdf
											_snake_case: Optional[int]	       =		mask * prev_known_part + (1.0 - mask) * prev_unknown_part

											if not return_dict:
																	return (
																	    pred_prev_sample,
																	    pred_original_sample,
																	)

											return RePaintSchedulerOutput(prev_sample=__snake_case							,  pred_original_sample=__snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	Dict							,  __snake_case						:	Optional[int]							,  __snake_case						:	List[str]							,  __snake_case						:	Dict=None							):





											'''simple docstring'''





											_snake_case: Optional[int]	       =		self.config.num_train_timesteps // self.num_inference_steps

											for i in range(__snake_case							):
																	_snake_case: Union[str, Any]	       =		self.betas[timestep + i]
																	if sample.device.type == "mps":
																							# randn does not work reproducibly on mps
																							_snake_case: Dict	       =		randn_tensor(sample.shape							,  dtype=sample.dtype							,  generator=__snake_case							)
																							_snake_case: Optional[int]	       =		noise.to(sample.device							)
																	else:
																							_snake_case: int	       =		randn_tensor(sample.shape							,  generator=__snake_case							,  device=sample.device							,  dtype=sample.dtype							)

																	# 10. Algorithm 1 Line 10 https://arxiv.org/pdf/2201.09865.pdf
																	_snake_case: int	       =		(1 - beta) ** 0.5 * sample + beta**0.5 * noise

											return sample


					def  SCREAMING_SNAKE_CASE_		(       self						:	str							,  __snake_case						:	torch.FloatTensor							,  __snake_case						:	torch.FloatTensor							,  __snake_case						:	torch.IntTensor							,  ):





											'''simple docstring'''





											raise NotImplementedError('Use `DDPMScheduler.add_noise()` to train for sampling with RePaint.'							)




					def __len__(       self						:	Dict							):





											'''simple docstring'''





											return self.config.num_train_timesteps












'''simple docstring'''






from __future__ import annotations

import copy
import os
import tempfile
from importlib import import_module
from math import isnan

from transformers import is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import _tf_gpu_memory_limit, require_tf, slow

from ..test_modeling_tf_common import ids_tensor


if is_tf_available():
		import numpy as np
		import tensorflow as tf

		from transformers import (
		    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
		    TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
		    TF_MODEL_FOR_MASKED_LM_MAPPING,
		    TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
		    TF_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING,
		    TF_MODEL_FOR_PRETRAINING_MAPPING,
		    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
		    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
		    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
		    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
		    TFSharedEmbeddings,
		)

		if _tf_gpu_memory_limit is not None:
				A :			Tuple							   =      tf.config.list_physical_devices('GPU')
				for gpu in gpus:
						# Restrict TensorFlow to only allocate x GB of memory on the GPUs
						try:
								tf.config.set_logical_device_configuration(
								    gpu, [tf.config.LogicalDeviceConfiguration(memory_limit=_tf_gpu_memory_limit)]
								)
								A :			Tuple							   =      tf.config.list_logical_devices('GPU')
								print('Logical GPUs', logical_gpus)
						except RuntimeError as e:
								# Virtual devices must be set before GPUs have been initialized
								print(e)







@require_tf
class       lowerCamelCase       :
					_SCREAMING_SNAKE_CASE				 =     None
					_SCREAMING_SNAKE_CASE				 =     ()
					_SCREAMING_SNAKE_CASE				 =     ()
					_SCREAMING_SNAKE_CASE				 =     True
					_SCREAMING_SNAKE_CASE				 =     True
					_SCREAMING_SNAKE_CASE				 =     True
					_SCREAMING_SNAKE_CASE				 =     False


					def  SCREAMING_SNAKE_CASE_		(       self						:	Union[str, Any]							,  __snake_case						:	Optional[Any]							,  __snake_case						:	int							,  __snake_case						:	Tuple=False							):





											'''simple docstring'''





											_snake_case: Tuple	       =		copy.deepcopy(__snake_case							)

											if model_class in get_values(__snake_case							):
																	_snake_case: Dict	       =		{
																	    k: tf.tile(tf.expand_dims(__snake_case							,  1							)							,  (1, self.model_tester.num_choices) + (1,) * (v.ndim - 1)							)
																	    if isinstance(__snake_case							,  tf.Tensor							) and v.ndim > 0
																	    else v
																	    for k, v in inputs_dict.items()
																	}

											if return_labels:
																	if model_class in get_values(__snake_case							):
																							_snake_case: Any	       =		tf.ones(self.model_tester.batch_size							,  dtype=tf.intaa							)
																	elif model_class in get_values(__snake_case							):
																							_snake_case: Dict	       =		tf.zeros(self.model_tester.batch_size							,  dtype=tf.intaa							)
																							_snake_case: Optional[Any]	       =		tf.zeros(self.model_tester.batch_size							,  dtype=tf.intaa							)
																	elif model_class in [
																	    *get_values(__snake_case							),
																	    *get_values(__snake_case							),
																	]:
																							_snake_case: int	       =		tf.zeros(self.model_tester.batch_size							,  dtype=tf.intaa							)
																	elif model_class in get_values(__snake_case							):
																							_snake_case: Optional[int]	       =		tf.zeros(self.model_tester.batch_size							,  dtype=tf.intaa							)
																	elif model_class in [
																	    *get_values(__snake_case							),
																	    *get_values(__snake_case							),
																	    *get_values(__snake_case							),
																	    *get_values(__snake_case							),
																	    *get_values(__snake_case							),
																	]:
																							_snake_case: Union[str, Any]	       =		tf.zeros(
																							    (self.model_tester.batch_size, self.model_tester.seq_length)							,  dtype=tf.intaa							)
											return inputs_dict


					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	Tuple							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: Optional[int]	       =		self.model_tester.prepare_config_and_inputs_for_common()
											for model_class in self.all_model_classes[:2]:
																	_snake_case: Tuple	       =		self._prepare_for_class(__snake_case							,  __snake_case							)
																	_snake_case: Dict	       =		model_class(__snake_case							)

																	@tf.function
																	def run_in_graph_mode():
																							return model(__snake_case							)

																	_snake_case: Dict	       =		run_in_graph_mode()
																	self.assertIsNotNone(__snake_case							)


					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[Any]							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: Optional[int]	       =		self.model_tester.prepare_config_and_inputs_for_common()
											for model_class in self.all_model_classes[:2]:
																	_snake_case: Optional[Any]	       =		self._prepare_for_class(__snake_case							,  __snake_case							)
																	_snake_case: str	       =		model_class(__snake_case							)

																	@tf.function(experimental_compile=__snake_case							)
																	def run_in_graph_mode():
																							return model(__snake_case							)

																	_snake_case: str	       =		run_in_graph_mode()
																	self.assertIsNotNone(__snake_case							)


					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[Any]							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: List[Any]	       =		self.model_tester.prepare_config_and_inputs_for_common()
											for model_class in self.all_model_classes[:2]:
																	_snake_case: Dict	       =		model_class(__snake_case							)
																	if getattr(__snake_case							,  'hf_compute_loss'							,  __snake_case							):
																							# Test that model correctly compute the loss with kwargs
																							_snake_case: Tuple	       =		self._prepare_for_class(inputs_dict.copy()							,  __snake_case							,  return_labels=__snake_case							)
																							# Is there a better way to remove these decoder inputs?
																							_snake_case: Union[str, Any]	       =		{
																							    key: val
																							    for key, val in prepared_for_class.items()
																							    if key not in ('head_mask', 'decoder_head_mask', 'cross_attn_head_mask', 'decoder_input_ids')
																							}

																							_snake_case: List[str]	       =		{
																							    'labels',
																							    'label',
																							    'label_ids',
																							    'start_positions',
																							    'start_position',
																							    'end_positions',
																							    'end_position',
																							    'next_sentence_label',
																							}
																							_snake_case: Tuple	       =		possible_label_cols.intersection(set(__snake_case							)							)
																							self.assertGreater(len(__snake_case							)							,  0							,  msg='No matching label names found!'							)
																							_snake_case: Optional[Any]	       =		{key: val for key, val in prepared_for_class.items() if key in label_names}
																							_snake_case: Union[str, Any]	       =		{key: val for key, val in prepared_for_class.items() if key not in label_names}
																							self.assertGreater(len(__snake_case							)							,  0							)

																							# Make sure it works with XLA!
																							model.compile(optimizer=tf.keras.optimizers.SGD(0.0							)							,  jit_compile=__snake_case							)
																							# Make sure the model fits without crashing regardless of where we pass the labels
																							_snake_case: Optional[Any]	       =		model.fit(
																							    __snake_case							,  validation_data=__snake_case							,  steps_per_epoch=1							,  validation_steps=1							,  shuffle=__snake_case							,  verbose=0							,  )
																							_snake_case: str	       =		history.history['loss'][0]
																							self.assertTrue(not isnan(__snake_case							)							)
																							_snake_case: Tuple	       =		history.history['val_loss'][0]
																							self.assertTrue(not isnan(__snake_case							)							)

																							# Now test it with separate labels, to make sure that path works in XLA too.
																							_snake_case: Any	       =		model_class(__snake_case							)
																							model.compile(optimizer=tf.keras.optimizers.SGD(0.0							)							,  jit_compile=__snake_case							)
																							_snake_case: str	       =		model.fit(
																							    __snake_case							,  __snake_case							,  validation_data=(inputs_minus_labels, labels)							,  steps_per_epoch=1							,  validation_steps=1							,  shuffle=__snake_case							,  verbose=0							,  )

																							_snake_case: List[str]	       =		history.history['loss'][0]
																							self.assertTrue(not isnan(__snake_case							)							)
																							_snake_case: Optional[Any]	       =		history.history['val_loss'][0]
																							self.assertTrue(not isnan(__snake_case							)							)


					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	List[Any]							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: List[Any]	       =		self.model_tester.prepare_config_and_inputs_for_common()
											_snake_case: int	       =		True
											_snake_case: int	       =		True

											if hasattr(__snake_case							,  'use_cache'							):
																	_snake_case: Optional[Any]	       =		True

											_snake_case: Union[str, Any]	       =		getattr(self.model_tester							,  'encoder_seq_length'							,  self.model_tester.seq_length							)
											_snake_case: int	       =		getattr(self.model_tester							,  'key_length'							,  __snake_case							)

											for model_class in self.all_model_classes[:2]:
																	_snake_case: Tuple	       =		self._prepare_for_class(__snake_case							,  __snake_case							)
																	_snake_case: int	       =		model_class(__snake_case							)
																	model.build()
																	_snake_case: str	       =		len(model(__snake_case							)							)

																	for key in list(class_inputs_dict.keys()							):
																							# Remove keys not in the serving signature, as the SavedModel will not be compiled to deal with them
																							if key not in model.input_signature:
																													del class_inputs_dict[key]
																							# Check it's a tensor, in case the inputs dict has some bools in it too
																							elif isinstance(class_inputs_dict[key]							,  tf.Tensor							) and class_inputs_dict[key].dtype.is_integer:
																													_snake_case: Dict	       =		tf.cast(class_inputs_dict[key]							,  tf.intaa							)

																	if set(class_inputs_dict.keys()							) != set(model.input_signature.keys()							):
																							continue  # Some models have inputs that the preparation functions don't create, we skip those

																	with tempfile.TemporaryDirectory() as tmpdirname:
																							model.save_pretrained(__snake_case							,  saved_model=__snake_case							)
																							_snake_case: int	       =		os.path.join(__snake_case							,  'saved_model'							,  '1'							)
																							_snake_case: Optional[Any]	       =		tf.keras.models.load_model(__snake_case							)
																							_snake_case: Any	       =		model(__snake_case							)

																							if self.is_encoder_decoder:
																													_snake_case: str	       =		outputs['encoder_hidden_states']
																													_snake_case: Union[str, Any]	       =		outputs['encoder_attentions']
																							else:
																													_snake_case: Optional[int]	       =		outputs['hidden_states']
																													_snake_case: Dict	       =		outputs['attentions']

																							self.assertEqual(len(__snake_case							)							,  __snake_case							)

																							_snake_case: int	       =		getattr(
																							    self.model_tester							,  'expected_num_hidden_layers'							,  self.model_tester.num_hidden_layers + 1							)

																							self.assertEqual(len(__snake_case							)							,  __snake_case							)
																							self.assertListEqual(
																							    list(output_hidden_states[0].shape[-2:]							)							,  [self.model_tester.seq_length, self.model_tester.hidden_size]							,  )

																							self.assertEqual(len(__snake_case							)							,  self.model_tester.num_hidden_layers							)
																							self.assertListEqual(
																							    list(output_attentions[0].shape[-3:]							)							,  [self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length]							,  )


					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	str							):





											'''simple docstring'''





											tf.keras.mixed_precision.set_global_policy('mixed_float16'							)

											# try/finally block to ensure subsequent tests run in float32
											try:
																	_snake_case       ,    _snake_case: List[str]	       =		self.model_tester.prepare_config_and_inputs_for_common()
																	for model_class in self.all_model_classes[:2]:
																							_snake_case: Dict	       =		self._prepare_for_class(__snake_case							,  __snake_case							)
																							_snake_case: Any	       =		model_class(__snake_case							)
																							_snake_case: List[str]	       =		model(__snake_case							)

																							self.assertIsNotNone(__snake_case							)
											finally:
																	tf.keras.mixed_precision.set_global_policy('float32'							)


					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	Dict							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: Optional[int]	       =		self.model_tester.prepare_config_and_inputs_for_common()
											# head_mask and decoder_head_mask has different shapes than other input args
											if "head_mask" in inputs_dict:
																	del inputs_dict["head_mask"]
											if "decoder_head_mask" in inputs_dict:
																	del inputs_dict["decoder_head_mask"]
											if "cross_attn_head_mask" in inputs_dict:
																	del inputs_dict["cross_attn_head_mask"]
											_snake_case: List[str]	       =		{
											    module_member
											    for model_class in self.all_model_classes
											    for module in (import_module(model_class.__module__							),)
											    for module_member_name in dir(__snake_case							)
											    if module_member_name.endswith('MainLayer'							)
											    for module_member in (getattr(__snake_case							,  __snake_case							),)
											    if isinstance(__snake_case							,  __snake_case							)
											    and tf.keras.layers.Layer in module_member.__bases__
											    and getattr(__snake_case							,  '_keras_serializable'							,  __snake_case							)
											}

											for main_layer_class in tf_main_layer_classes:
																	# T5MainLayer needs an embed_tokens parameter when called without the inputs_embeds parameter
																	if "T5" in main_layer_class.__name__:
																							# Take the same values than in TFT5ModelTester for this shared layer
																							_snake_case: Optional[Any]	       =		TFSharedEmbeddings(self.model_tester.vocab_size							,  self.model_tester.hidden_size							,  name='shared'							)
																							_snake_case: str	       =		False
																							_snake_case: Tuple	       =		main_layer_class(__snake_case							,  embed_tokens=__snake_case							)
																	else:
																							_snake_case: List[Any]	       =		main_layer_class(__snake_case							)

																	_snake_case: Tuple	       =		{
																	    name: tf.keras.Input(tensor.shape[1:]							,  dtype=tensor.dtype							) for name, tensor in inputs_dict.items()
																	}

																	if hasattr(self.model_tester							,  'num_labels'							):
																							_snake_case: Dict	       =		self.model_tester.num_labels
																	else:
																							_snake_case: Tuple	       =		2

																	_snake_case: Tuple	       =		tf.data.Dataset.from_tensor_slices(
																	    (inputs_dict, np.ones((self.model_tester.batch_size, self.model_tester.seq_length, num_labels, 1)							))							).batch(1							)

																	_snake_case: List[str]	       =		main_layer(__snake_case							)[0]
																	_snake_case: List[Any]	       =		tf.keras.layers.Dense(__snake_case							,  activation='softmax'							,  name='outputs'							)(__snake_case							)
																	_snake_case: List[str]	       =		tf.keras.models.Model(inputs=__snake_case							,  outputs=[outputs]							)

																	model.compile(loss='binary_crossentropy'							,  optimizer='adam'							,  metrics=['binary_accuracy']							)
																	model.fit(__snake_case							,  epochs=1							)

																	with tempfile.TemporaryDirectory() as tmpdirname:
																							_snake_case: Tuple	       =		os.path.join(__snake_case							,  'keras_model.h5'							)
																							model.save(__snake_case							)
																							if "T5" in main_layer_class.__name__:
																													_snake_case: List[str]	       =		tf.keras.models.load_model(
																													    __snake_case							,  custom_objects={
																													        main_layer_class.__name__: main_layer_class,
																													        'TFSharedEmbeddings': TFSharedEmbeddings,
																													    }							,  )
																							else:
																													_snake_case: Optional[int]	       =		tf.keras.models.load_model(
																													    __snake_case							,  custom_objects={main_layer_class.__name__: main_layer_class}							)
																							assert isinstance(__snake_case							,  tf.keras.Model							)
																							model(__snake_case							)


					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	Tuple							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: Optional[Any]	       =		self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes[:2]:
																	_snake_case: List[Any]	       =		model_class(__snake_case							)

																	_snake_case: Optional[int]	       =		copy.deepcopy(__snake_case							)

																	if not self.is_encoder_decoder:
																							_snake_case: Dict	       =		inputs['input_ids']
																							del inputs["input_ids"]
																	else:
																							_snake_case: Any	       =		inputs['input_ids']
																							_snake_case: List[str]	       =		inputs.get('decoder_input_ids'							,  __snake_case							)
																							del inputs["input_ids"]
																							inputs.pop('decoder_input_ids'							,  __snake_case							)

																	if not self.is_encoder_decoder:
																							_snake_case: str	       =		model.get_input_embeddings()(__snake_case							)
																	else:
																							_snake_case: Optional[int]	       =		model.get_input_embeddings()(__snake_case							)
																							_snake_case: List[str]	       =		model.get_input_embeddings()(__snake_case							)

																	_snake_case: str	       =		self._prepare_for_class(__snake_case							,  __snake_case							)

																	@tf.function
																	def run_in_graph_mode():
																							return model(__snake_case							)

																	_snake_case: Optional[Any]	       =		run_in_graph_mode()
																	self.assertIsNotNone(__snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	Dict							,  __snake_case						:	Optional[Any]							,  __snake_case						:	List[Any]							):





											'''simple docstring'''





											_snake_case: int	       =		[]
											if model.config.bos_token_id is not None:
																	special_tokens.append(model.config.bos_token_id							)
											if model.config.pad_token_id is not None:
																	special_tokens.append(model.config.pad_token_id							)
											if model.config.eos_token_id is not None:
																	special_tokens.append(model.config.eos_token_id							)

											# create random bad tokens that are not special tokens
											_snake_case: Tuple	       =		[]
											while len(__snake_case							) < num_bad_tokens:
																	_snake_case: Optional[Any]	       =		tf.squeeze(ids_tensor((1, 1)							,  self.model_tester.vocab_size							)							,  0							).numpy()[0]
																	if token not in special_tokens:
																							bad_tokens.append(__snake_case							)
											return bad_tokens


					def  SCREAMING_SNAKE_CASE_		(       self						:	Any							,  __snake_case						:	Tuple							):





											'''simple docstring'''





											for token_id in output_ids[0].numpy().tolist():
																	self.assertGreaterEqual(__snake_case							,  0							)
																	self.assertLess(__snake_case							,  self.model_tester.vocab_size							)




					def  SCREAMING_SNAKE_CASE_		(       self						:	Union[str, Any]							,  __snake_case						:	List[str]							,  __snake_case						:	int							):





											'''simple docstring'''





											for bad_word_ids in bad_words_ids:
																	# for all slices in batch
																	for generated_ids_slice in generated_ids:
																							# for all word idx
																							for i in range(len(__snake_case							)							,  len(__snake_case							)							):
																													# if tokens match
																													if generated_ids_slice[i - len(__snake_case							) : i] == bad_word_ids:
																																			return True
											return False





