

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:
														lowerCAmelCase__             =		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)]
																												)
																												lowerCAmelCase__             =		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 						lowercase     :


		"""simple docstring"""




		a__       							=	None
		a__       							=	()
		a__       							=	()
		a__       							=	True
		a__       							=	True
		a__       							=	True
		a__       							=	False


		def 		A__       (					self						,							__snake_case						,							__snake_case						,							__snake_case=False):
			_UpperCamelCase :		List[str]								=    copy.deepcopy(__snake_case)

			if model_class in get_values(__snake_case):
				_UpperCamelCase :		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):
					_UpperCamelCase :		Union[str, Any]								=    tf.ones(self.model_tester.batch_size						,							dtype=tf.intaa)
				elif model_class in get_values(__snake_case):
					_UpperCamelCase :		Optional[Any]								=    tf.zeros(self.model_tester.batch_size						,							dtype=tf.intaa)
					_UpperCamelCase :		Tuple								=    tf.zeros(self.model_tester.batch_size						,							dtype=tf.intaa)
				elif model_class in [
				    *get_values(__snake_case),
				    *get_values(__snake_case),
				]:
					_UpperCamelCase :		Optional[int]								=    tf.zeros(self.model_tester.batch_size						,							dtype=tf.intaa)
				elif model_class in get_values(__snake_case):
					_UpperCamelCase :		Dict								=    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),
				]:
					_UpperCamelCase :		Tuple								=    tf.zeros(
					    (self.model_tester.batch_size, self.model_tester.seq_length)						,							dtype=tf.intaa)
			return inputs_dict


		@slow
		def 		A__       (					self):
			_UpperCamelCase						,		_UpperCamelCase :		Optional[Any]								=    self.model_tester.prepare_config_and_inputs_for_common()
			for model_class in self.all_model_classes[:2]:
				_UpperCamelCase :		Dict								=    self._prepare_for_class(__snake_case						,							__snake_case)
				_UpperCamelCase :		str								=    model_class(__snake_case)

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

				_UpperCamelCase :		List[str]								=    run_in_graph_mode()
				self.assertIsNotNone(__snake_case)


		@slow
		def 		A__       (					self):
			_UpperCamelCase						,		_UpperCamelCase :		Tuple								=    self.model_tester.prepare_config_and_inputs_for_common()
			for model_class in self.all_model_classes[:2]:
				_UpperCamelCase :		Dict								=    self._prepare_for_class(__snake_case						,							__snake_case)
				_UpperCamelCase :		Union[str, Any]								=    model_class(__snake_case)

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

				_UpperCamelCase :		List[str]								=    run_in_graph_mode()
				self.assertIsNotNone(__snake_case)


		@slow
		def 		A__       (					self):
			# This is a copy of the test_keras_fit method, but we use XLA compilation instead of eager
			_UpperCamelCase						,		_UpperCamelCase :		Tuple								=    self.model_tester.prepare_config_and_inputs_for_common()
			for model_class in self.all_model_classes[:2]:
				_UpperCamelCase :		Optional[int]								=    model_class(__snake_case)
				if getattr(__snake_case						,							'hf_compute_loss'						,							__snake_case):
					# Test that model correctly compute the loss with kwargs
					_UpperCamelCase :		Tuple								=    self._prepare_for_class(inputs_dict.copy()						,							__snake_case						,							return_labels=__snake_case)
					# Is there a better way to remove these decoder inputs?
					_UpperCamelCase :		int								=    {
					    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')
					}

					_UpperCamelCase :		Any								=    {
					    'labels',
					    'label',
					    'label_ids',
					    'start_positions',
					    'start_position',
					    'end_positions',
					    'end_position',
					    'next_sentence_label',
					}
					_UpperCamelCase :		Any								=    possible_label_cols.intersection(set(__snake_case))
					self.assertGreater(len(__snake_case)						,							0						,							msg='No matching label names found!')
					_UpperCamelCase :		Tuple								=    {key: val for key, val in prepared_for_class.items() if key in label_names}
					_UpperCamelCase :		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
					_UpperCamelCase :		int								=    model.fit(
					    __snake_case						,							validation_data=__snake_case						,							steps_per_epoch=1						,							validation_steps=1						,							shuffle=__snake_case						,							verbose=0						,							)
					_UpperCamelCase :		Optional[Any]								=    history.history['loss'][0]
					self.assertTrue(not isnan(__snake_case))
					_UpperCamelCase :		List[Any]								=    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.
					_UpperCamelCase :		str								=    model_class(__snake_case)
					model.compile(optimizer=tf.keras.optimizers.SGD(0.0)						,							jit_compile=__snake_case)
					_UpperCamelCase :		List[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						,							)

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


		@slow
		def 		A__       (					self):
			_UpperCamelCase						,		_UpperCamelCase :		List[Any]								=    self.model_tester.prepare_config_and_inputs_for_common()
			_UpperCamelCase :		Dict								=    True
			_UpperCamelCase :		List[str]								=    True

			if hasattr(__snake_case						,							'use_cache'):
				_UpperCamelCase :		Any								=    True

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

			for model_class in self.all_model_classes[:2]:
				_UpperCamelCase :		Union[str, Any]								=    self._prepare_for_class(__snake_case						,							__snake_case)
				_UpperCamelCase :		Optional[Any]								=    model_class(__snake_case)
				model.build()
				_UpperCamelCase :		Tuple								=    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:
						_UpperCamelCase :		Optional[Any]								=    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)
					_UpperCamelCase :		Optional[Any]								=    os.path.join(__snake_case						,							'saved_model'						,							'1')
					_UpperCamelCase :		Optional[Any]								=    tf.keras.models.load_model(__snake_case)
					_UpperCamelCase :		List[str]								=    model(__snake_case)

					if self.is_encoder_decoder:
						_UpperCamelCase :		str								=    outputs['encoder_hidden_states']
						_UpperCamelCase :		List[Any]								=    outputs['encoder_attentions']
					else:
						_UpperCamelCase :		Tuple								=    outputs['hidden_states']
						_UpperCamelCase :		Union[str, Any]								=    outputs['attentions']

					self.assertEqual(len(__snake_case)						,							__snake_case)

					_UpperCamelCase :		Union[str, Any]								=    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 		A__       (					self):
			tf.keras.mixed_precision.set_global_policy('mixed_float16')

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

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


		@slow
		def 		A__       (					self):
			_UpperCamelCase						,		_UpperCamelCase :		List[Any]								=    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"]
			_UpperCamelCase :		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
					_UpperCamelCase :		List[str]								=    TFSharedEmbeddings(self.model_tester.vocab_size						,							self.model_tester.hidden_size						,							name='shared')
					_UpperCamelCase :		List[Any]								=    False
					_UpperCamelCase :		Union[str, Any]								=    main_layer_class(__snake_case						,							embed_tokens=__snake_case)
				else:
					_UpperCamelCase :		Optional[int]								=    main_layer_class(__snake_case)

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

				if hasattr(self.model_tester						,							'num_labels'):
					_UpperCamelCase :		int								=    self.model_tester.num_labels
				else:
					_UpperCamelCase :		Union[str, Any]								=    2

				_UpperCamelCase :		Dict								=    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)

				_UpperCamelCase :		Any								=    main_layer(__snake_case)[0]
				_UpperCamelCase :		Optional[Any]								=    tf.keras.layers.Dense(__snake_case						,							activation='softmax'						,							name='outputs')(__snake_case)
				_UpperCamelCase :		Dict								=    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:
					_UpperCamelCase :		Union[str, Any]								=    os.path.join(__snake_case						,							'keras_model.h5')
					model.save(__snake_case)
					if "T5" in main_layer_class.__name__:
						_UpperCamelCase :		List[str]								=    tf.keras.models.load_model(
						    __snake_case						,							custom_objects={
						        main_layer_class.__name__: main_layer_class,
						        'TFSharedEmbeddings': TFSharedEmbeddings,
						    }						,							)
					else:
						_UpperCamelCase :		List[str]								=    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 		A__       (					self):
			_UpperCamelCase						,		_UpperCamelCase :		str								=    self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes[:2]:
				_UpperCamelCase :		Optional[Any]								=    model_class(__snake_case)

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

				if not self.is_encoder_decoder:
					_UpperCamelCase :		List[Any]								=    inputs['input_ids']
					del inputs["input_ids"]
				else:
					_UpperCamelCase :		List[Any]								=    inputs['input_ids']
					_UpperCamelCase :		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:
					_UpperCamelCase :		List[Any]								=    model.get_input_embeddings()(__snake_case)
				else:
					_UpperCamelCase :		Tuple								=    model.get_input_embeddings()(__snake_case)
					_UpperCamelCase :		List[str]								=    model.get_input_embeddings()(__snake_case)

				_UpperCamelCase :		Union[str, Any]								=    self._prepare_for_class(__snake_case						,							__snake_case)

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

				_UpperCamelCase :		List[Any]								=    run_in_graph_mode()
				self.assertIsNotNone(__snake_case)


		def 		A__       (					self						,							__snake_case						,							__snake_case):
			# special tokens cannot be bad tokens
			_UpperCamelCase :		str								=    []
			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
			_UpperCamelCase :		Optional[int]								=    []
			while len(__snake_case) < num_bad_tokens:
				_UpperCamelCase :		Tuple								=    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 		A__       (					self						,							__snake_case):
			for token_id in output_ids[0].numpy().tolist():
				self.assertGreaterEqual(__snake_case						,							0)
				self.assertLess(__snake_case						,							self.model_tester.vocab_size)






		def 		A__       (					self						,							__snake_case						,							__snake_case):
			# for all bad word tokens
			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







import os

from apache_beam.io.filesystems import FileSystems
from apache_beam.pipeline import Pipeline

from .logging import get_logger


lowerCAmelCase__             =		2 << 2_0  # 2mb
lowerCAmelCase__             =		get_logger(__name__)







class 						lowercase     (					_lowercase				):


		"""simple docstring"""




		def 		A__       (					self):
			_UpperCamelCase :		Optional[Any]								=    self._options.get_all_options().get('runner')
			return runner in [None, "DirectRunner", "PortableRunner"]






def        lowerCamelCase_  (				UpperCAmelCase_				:  Any				,							UpperCAmelCase_				:  Optional[Any]				,							UpperCAmelCase_				:  str=False						)    ->					str:




	'''simple docstring'''




	_UpperCamelCase :		Optional[int]								=    FileSystems
	if fs.exists(UpperCAmelCase_						):
		if force_upload:
			logger.info(F'''Remote path already exist: {remote_file_path}. Overwriting it as force_upload=True.'''						)
		else:
			logger.info(F'''Remote path already exist: {remote_file_path}. Skipping it as force_upload=False.'''						)
			return
	with fs.create(UpperCAmelCase_						) as remote_file:
		with open(UpperCAmelCase_				,							'rb'						) as local_file:
			_UpperCamelCase :		List[str]								=    local_file.read(UpperCAmelCase_						)
			while chunk:
				remote_file.write(UpperCAmelCase_						)
				_UpperCamelCase :		List[str]								=    local_file.read(UpperCAmelCase_						)







def        lowerCamelCase_  (				UpperCAmelCase_				:  Any				,							UpperCAmelCase_				:  List[Any]				,							UpperCAmelCase_				:  Optional[int]=False						)    ->					Any:




	'''simple docstring'''




	_UpperCamelCase :		List[str]								=    FileSystems
	if os.path.exists(UpperCAmelCase_						):
		if force_download:
			logger.info(F'''Local path already exist: {remote_file_path}. Overwriting it as force_upload=True.'''						)
		else:
			logger.info(F'''Local path already exist: {remote_file_path}. Skipping it as force_upload=False.'''						)
			return
	with fs.open(UpperCAmelCase_						) as remote_file:
		with open(UpperCAmelCase_				,							'wb'						) as local_file:
			_UpperCamelCase :		Union[str, Any]								=    remote_file.read(UpperCAmelCase_						)
			while chunk:
				local_file.write(UpperCAmelCase_						)
				_UpperCamelCase :		List[str]								=    remote_file.read(UpperCAmelCase_						)




