





'''simple docstring'''






import csv
from collections import defaultdict
from dataclasses import dataclass, field
from typing import List, Optional

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import ScalarFormatter

from transformers import HfArgumentParser





def    lowercase_   (		lowercase__=None					,	lowercase__=None    ) ->int:
						return field(default_factory=lambda: default					,	metadata=lowercase__    )







@dataclass
class       lowerCamelCase       :
					_SCREAMING_SNAKE_CASE				 =     field(
					    metadata={"help": "The csv file to plot."}  ,				)
					_SCREAMING_SNAKE_CASE				 =     field(
					    default=__UpperCAmelCase  ,				metadata={"help": "Whether to plot along batch size or sequence length. Defaults to sequence length."}  ,				)
					_SCREAMING_SNAKE_CASE				 =     field(
					    default=__UpperCAmelCase  ,				metadata={"help": "Whether the csv file has time results or memory results. Defaults to memory results."}  ,				)
					_SCREAMING_SNAKE_CASE				 =     field(
					    default=__UpperCAmelCase  ,				metadata={"help": "Disable logarithmic scale when plotting"}  ,				)
					_SCREAMING_SNAKE_CASE				 =     field(
					    default=__UpperCAmelCase  ,				metadata={
					        "help": "Whether the csv file has training results or inference results. Defaults to inference results."
					    }  ,				)
					_SCREAMING_SNAKE_CASE				 =     field(
					    default=__UpperCAmelCase  ,				metadata={"help": "Filename under which the plot will be saved. If unused no plot is saved."}  ,				)
					_SCREAMING_SNAKE_CASE				 =     list_field(
					    default=__UpperCAmelCase  ,				metadata={"help": "List of model names that are used instead of the ones in the csv file."}	)





def    lowercase_   (		lowercase__    ) ->List[str]:
						try:
												int(lowercase__    )
												return True
						except ValueError:
												return False





def    lowercase_   (		lowercase__    ) ->List[Any]:
						try:
												float(lowercase__    )
												return True
						except ValueError:
												return False







class       lowerCamelCase       :


					def __init__(       self						:	List[str]							,  __snake_case						:	str							):





											'''simple docstring'''





											_snake_case: Tuple	       =		args
											_snake_case: List[Any]	       =		defaultdict(lambda: {"bsz": [], "seq_len": [], "result": {}}							)

											with open(self.args.csv_file							,  newline=''							) as csv_file:
																	_snake_case: Optional[int]	       =		csv.DictReader(__snake_case							)
																	for row in reader:
																							_snake_case: Union[str, Any]	       =		row['model']
																							self.result_dict[model_name]["bsz"].append(int(row['batch_size']							)							)
																							self.result_dict[model_name]["seq_len"].append(int(row['sequence_length']							)							)
																							if can_convert_to_int(row['result']							):
																													# value is not None
																													_snake_case: Optional[int]	       =		int(row['result']							)
																							elif can_convert_to_float(row['result']							):
																													# value is not None
																													_snake_case: Dict	       =		float(row['result']							)




					def  SCREAMING_SNAKE_CASE_		(       self						:	List[Any]							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: str	       =		plt.subplots()
											_snake_case: Tuple	       =		'Time usage' if self.args.is_time else 'Memory usage'
											_snake_case: Dict	       =		title_str + ' for training' if self.args.is_train else title_str + ' for inference'

											if not self.args.no_log_scale:
																	# set logarithm scales
																	ax.set_xscale('log'							)
																	ax.set_yscale('log'							)

											for axis in [ax.xaxis, ax.yaxis]:
																	axis.set_major_formatter(ScalarFormatter()							)

											for model_name_idx, model_name in enumerate(self.result_dict.keys()							):
																	_snake_case: Any	       =		sorted(set(self.result_dict[model_name]['bsz']							)							)
																	_snake_case: Tuple	       =		sorted(set(self.result_dict[model_name]['seq_len']							)							)
																	_snake_case: List[str]	       =		self.result_dict[model_name]['result']

																	((_snake_case)       ,    (_snake_case)): Union[str, Any]	       =		(
																	    (batch_sizes, sequence_lengths) if self.args.plot_along_batch else (sequence_lengths, batch_sizes)
																	)

																	_snake_case: List[str]	       =		(
																	    model_name if self.args.short_model_names is None else self.args.short_model_names[model_name_idx]
																	)

																	for inner_loop_value in inner_loop_array:
																							if self.args.plot_along_batch:
																													_snake_case: Union[str, Any]	       =		np.asarray(
																													    [results[(x, inner_loop_value)] for x in x_axis_array if (x, inner_loop_value) in results]							,  dtype=__snake_case							,  )
																							else:
																													_snake_case: Optional[int]	       =		np.asarray(
																													    [results[(inner_loop_value, x)] for x in x_axis_array if (inner_loop_value, x) in results]							,  dtype=np.floataa							,  )

																							((_snake_case)       ,    (_snake_case)): Any	       =		(
																							    ('batch_size', 'len') if self.args.plot_along_batch else ('in #tokens', 'bsz')
																							)

																							_snake_case: Dict	       =		np.asarray(__snake_case							,  __snake_case							)[: len(__snake_case							)]
																							plt.scatter(
																							    __snake_case							,  __snake_case							,  label=f'''{label_model_name} - {inner_loop_label}: {inner_loop_value}'''							)
																							plt.plot(__snake_case							,  __snake_case							,  '--'							)

																	title_str += f''' {label_model_name} vs.'''

											_snake_case: Dict	       =		title_str[:-4]
											_snake_case: List[str]	       =		'Time in s' if self.args.is_time else 'Memory in MB'

											# plot
											plt.title(__snake_case							)
											plt.xlabel(__snake_case							)
											plt.ylabel(__snake_case							)
											plt.legend()

											if self.args.figure_png_file is not None:
																	plt.savefig(self.args.figure_png_file							)
											else:
																	plt.show()





def    lowercase_   (		) ->List[Any]:
						_snake_case: Any	       =		HfArgumentParser(lowercase__    )
						_snake_case: str	       =		parser.parse_args_into_dataclasses()[0]
						_snake_case: List[str]	       =		Plot(args=lowercase__    )
						plot.plot()


if __name__ == "__main__":
		main()












'''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





