


'''simple docstring'''




import argparse
import logging
import os
import sys
import tempfile
from pathlib import Path

import lightning_base
import pytest
import pytorch_lightning as pl
import torch
from convert_pl_checkpoint_to_hf import convert_pl_to_hf
from distillation import distill_main
from finetune import SummarizationModule, main
from huggingface_hub import list_models
from parameterized import parameterized
from run_eval import generate_summaries_or_translations
from torch import nn

from transformers import AutoConfig, AutoModelForSeqaSeqLM
from transformers.testing_utils import CaptureStderr, CaptureStdout, TestCasePlus, require_torch_gpu, slow
from utils import label_smoothed_nll_loss, lmap, load_json


logging.basicConfig(level=logging.DEBUG)

_lowerCAmelCase	:		List[str]		  =						logging.getLogger()
_lowerCAmelCase	:		Tuple		  =						torch.cuda.is_available()
_lowerCAmelCase	:		Tuple		  =						{
    'max_tokens_per_batch': None,
    'supervise_forward': True,
    'normalize_hidden': True,
    'label_smoothing': 0.2,
    'eval_max_gen_length': None,
    'eval_beams': 1,
    'val_metric': 'loss',
    'save_top_k': 1,
    'adafactor': True,
    'early_stopping_patience': 2,
    'logger_name': 'default',
    'length_penalty': 0.5,
    'cache_dir': '',
    'task': 'summarization',
    'num_workers': 2,
    'alpha_hid': 0,
    'freeze_embeds': True,
    'enc_only': False,
    'tgt_suffix': '',
    'resume_from_checkpoint': None,
    'sortish_sampler': True,
    'student_decoder_layers': 1,
    'val_check_interval': 1.0,
    'output_dir': '',
    'fp16': False,  # TODO(SS): set this to CUDA_AVAILABLE if ci installs apex or start using native amp
    'no_teacher': False,
    'fp16_opt_level': 'O1',
    'gpus': 1 if CUDA_AVAILABLE else 0,
    'n_tpu_cores': 0,
    'max_grad_norm': 1.0,
    'do_train': True,
    'do_predict': True,
    'accumulate_grad_batches': 1,
    'server_ip': '',
    'server_port': '',
    'seed': 42,
    'model_name_or_path': 'sshleifer/bart-tiny-random',
    'config_name': '',
    'tokenizer_name': 'facebook/bart-large',
    'do_lower_case': False,
    'learning_rate': 0.3,
    'lr_scheduler': 'linear',
    'weight_decay': 0.0,
    'adam_epsilon': 1E-08,
    'warmup_steps': 0,
    'max_epochs': 1,
    'train_batch_size': 2,
    'eval_batch_size': 2,
    'max_source_length': 12,
    'max_target_length': 12,
    'val_max_target_length': 12,
    'test_max_target_length': 12,
    'fast_dev_run': False,
    'no_cache': False,
    'n_train': -1,
    'n_val': -1,
    'n_test': -1,
    'student_encoder_layers': 1,
    'freeze_encoder': False,
    'auto_scale_batch_size': False,
    'overwrite_output_dir': False,
    'student': None,
}

def       __UpperCamelCase     (     _A     :					Path       , _A     :					list					)     ->   Optional[int]:




	"""simple docstring"""





	lowerCAmelCase       :     Optional[int]   				=			'\n'.join(__snake_case					)
	Path(__snake_case					).open('w'					).writelines(__snake_case					)


_lowerCAmelCase	:		int		  =						[' Sam ate lunch today.', 'Sams lunch ingredients.']
_lowerCAmelCase	:		Optional[Any]		  =						['A very interesting story about what I ate for lunch.', 'Avocado, celery, turkey, coffee']
_lowerCAmelCase	:		Any		  =						'patrickvonplaten/t5-tiny-random'
_lowerCAmelCase	:		Union[str, Any]		  =						'sshleifer/t5-tinier-random'
_lowerCAmelCase	:		str		  =						'sshleifer/bart-tiny-random'
_lowerCAmelCase	:		int		  =						'sshleifer/tiny-mbart'
_lowerCAmelCase	:		List[str]		  =						'sshleifer/tiny-marian-en-de'
_lowerCAmelCase	:		Any		  =						'stas/tiny-wmt19-en-de'


_lowerCAmelCase	:		Union[str, Any]		  =						logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)
logging.disable(logging.CRITICAL)  # remove noisy download output from tracebacks

def       __UpperCamelCase     (     _A     :					List[Any]					)     ->   List[Any]:




	"""simple docstring"""





	for split in ["train", "val", "test"]:
		_dump_articles(os.path.join(__snake_case       , F"{split}.source"					)       , __snake_case					)
		_dump_articles(os.path.join(__snake_case       , F"{split}.target"					)       , __snake_case					)
	return tmp_dir



class 			lowerCAmelCase					(					a_				):



	@classmethod
	def 					lowercase  (       cls     ):
		logging.disable(logging.CRITICAL     )  # remove noisy download output from tracebacks
		return cls



	@slow
	@require_torch_gpu
	def 					lowercase  (       self     ):
		lowerCAmelCase       :     Union[str, Any]   				=			list_models()
		lowerCAmelCase       :     int   				=			'sshleifer'
		lowerCAmelCase       :     Optional[Any]   				=			[x.modelId for x in model_list if x.modelId.startswith(_lowerCamelCase     )]
		lowerCAmelCase       :     str   				=			['sshleifer/blenderbot-3B', 'sshleifer/blenderbot-90M']
		lowerCAmelCase       :     Tuple   				=			[]
		for m in model_ids:
			if m in allowed_to_be_broken:
				continue
			try:
				AutoConfig.from_pretrained(_lowerCamelCase     )
			except Exception:
				failures.append(_lowerCamelCase     )
		assert not failures, f"The following models could not be loaded through AutoConfig: {failures}"



	def 					lowercase  (       self     ):
		lowerCAmelCase       :     Dict   				=			{'student_encoder_layers': 2, 'student_decoder_layers': 1, 'no_teacher': True}
		self._test_distiller_cli(_lowerCamelCase     )



	def 					lowercase  (       self     ):
		lowerCAmelCase       :     Any   				=			{
		    'student_encoder_layers': 2,
		    'student_decoder_layers': 1,
		    'max_epochs': 4,
		    'val_check_interval': 0.2_5,
		    'alpha_hid': 2.0,
		    'model_name_or_path': 'IGNORE_THIS_IT_DOESNT_GET_USED',
		}
		lowerCAmelCase       :     Dict   				=			self._test_distiller_cli(_lowerCamelCase			,       check_contents=_lowerCamelCase     )

		lowerCAmelCase       :     Union[str, Any]   				=			list(Path(model.output_dir     ).glob('*.ckpt'     )     )
		self.assertEqual(1			,       len(_lowerCamelCase     )     )
		lowerCAmelCase       :     Union[str, Any]   				=			list(Path(model.output_dir     ).glob('**/*.bin'     )     )
		self.assertEqual(len(_lowerCamelCase     )			,       2     )
		lowerCAmelCase       :     int   				=			lmap(str.strip			,       Path(model.hparams.data_dir     ).joinpath('test.source'     ).open().readlines()     )
		lowerCAmelCase       :     Optional[int]   				=			tempfile.mktemp()  # XXX: not being cleaned up
		generate_summaries_or_translations(_lowerCamelCase			,       _lowerCamelCase			,       str(model.output_dir / 'best_tfmr'     )     )
		self.assertTrue(Path(_lowerCamelCase     ).exists()     )

		lowerCAmelCase       :     int   				=			self.get_auto_remove_tmp_dir()
		convert_pl_to_hf(ckpts[0]			,       transformer_ckpts[0].parent			,       _lowerCamelCase     )
		assert os.path.exists(os.path.join(_lowerCamelCase			,       'pytorch_model.bin'     )     )



	def 					lowercase  (       self     ):
		lowerCAmelCase       :     Optional[int]   				=			AutoModelForSeqaSeqLM.from_pretrained(_lowerCamelCase     )
		lowerCAmelCase				, lowerCAmelCase       :     int   				=			model.dummy_inputs['input_ids'], model.dummy_inputs['attention_mask']
		lowerCAmelCase       :     str   				=			torch.tensor([[0, 4, 8, 2], [0, 8, 2, 1]]			,       dtype=torch.long			,       device=model.device     )
		lowerCAmelCase       :     Any   				=			target_ids[:, :-1].contiguous()  # Why this line?
		lowerCAmelCase       :     Optional[int]   				=			target_ids[:, 1:].clone()  # why clone?
		lowerCAmelCase       :     str   				=			model(
		    _lowerCamelCase			,       attention_mask=_lowerCamelCase			,       decoder_input_ids=_lowerCamelCase			,       labels=_lowerCamelCase			,       use_cache=_lowerCamelCase     ).loss

		lowerCAmelCase       :     List[str]   				=			model(_lowerCamelCase			,       attention_mask=_lowerCamelCase			,       decoder_input_ids=_lowerCamelCase			,       use_cache=_lowerCamelCase     ).logits

		lowerCAmelCase       :     int   				=			nn.functional.log_softmax(_lowerCamelCase			,       dim=-1     )
		lowerCAmelCase				, lowerCAmelCase       :     Any   				=			label_smoothed_nll_loss(
		    _lowerCamelCase			,       _lowerCamelCase			,       0.1			,       ignore_index=model.config.pad_token_id     )
		with self.assertRaises(_lowerCamelCase     ):
			# TODO: understand why this breaks
			self.assertEqual(_lowerCamelCase			,       _lowerCamelCase     )



	def 					lowercase  (       self     ):
		lowerCAmelCase       :     Dict   				=			{
		    'student_encoder_layers': 2,
		    'student_decoder_layers': 1,
		    'num_train_epochs': 4,
		    'val_check_interval': 0.2_5,
		    'alpha_hid': 2.0,
		    'task': 'translation',
		    'model_name_or_path': 'IGNORE_THIS_IT_DOESNT_GET_USED',
		    'tokenizer_name': MBART_TINY,
		    'teacher': MBART_TINY,
		    'src_lang': 'en_XX',
		    'tgt_lang': 'ro_RO',
		}
		lowerCAmelCase       :     Dict   				=			self._test_distiller_cli(_lowerCamelCase			,       check_contents=_lowerCamelCase     )
		assert model.model.config.model_type == "mbart"

		lowerCAmelCase       :     Dict   				=			list(Path(model.output_dir     ).glob('*.ckpt'     )     )
		self.assertEqual(1			,       len(_lowerCamelCase     )     )
		lowerCAmelCase       :     int   				=			list(Path(model.output_dir     ).glob('**/*.bin'     )     )
		lowerCAmelCase       :     Dict   				=			list(Path(model.output_dir     ).glob('best_tfmr/*'     )     )
		assert len(_lowerCamelCase     ) > 2
		self.assertEqual(len(_lowerCamelCase     )			,       2     )



	def 					lowercase  (       self     ):
		lowerCAmelCase       :     List[str]   				=			{
		    'student_encoder_layers': 1,
		    'student_decoder_layers': 1,
		    'alpha_hid': 2.0,
		    'teacher': T5_TINY,
		    'model_name_or_path': T5_TINY,
		    'tokenizer_name': T5_TINY,
		}
		self._test_distiller_cli(_lowerCamelCase     )



	def 					lowercase  (       self     ):
		lowerCAmelCase       :     str   				=			{
		    'teacher': T5_TINY,
		    'student': T5_TINIER,
		    'model_name_or_path': T5_TINIER,
		    'tokenizer_name': T5_TINIER,
		}
		self._test_distiller_cli(_lowerCamelCase     )



	def 					lowercase  (       self			,       snake_case__			,       snake_case__=True     ):
		lowerCAmelCase       :     str   				=			{
		    'label_smoothing': 0.0,
		    'early_stopping_patience': -1,
		    'train_batch_size': 1,
		    'eval_batch_size': 2,
		    'max_epochs': 2,
		    'alpha_mlm': 0.2,
		    'alpha_ce': 0.8,
		    'do_predict': True,
		    'model_name_or_path': 'sshleifer/tinier_bart',
		    'teacher': CHEAP_ARGS['model_name_or_path'],
		    'val_check_interval': 0.5,
		}
		default_updates.update(_lowerCamelCase     )
		lowerCAmelCase       :     str   				=			CHEAP_ARGS.copy()
		lowerCAmelCase       :     int   				=			make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()     )
		lowerCAmelCase       :     List[str]   				=			self.get_auto_remove_tmp_dir()

		args_d.update(data_dir=_lowerCamelCase			,       output_dir=_lowerCamelCase			,       **_lowerCamelCase     )
		lowerCAmelCase       :     List[str]   				=			distill_main(argparse.Namespace(**_lowerCamelCase     )     )
		if not check_contents:
			return model
		lowerCAmelCase       :     Optional[int]   				=			os.listdir(_lowerCamelCase     )
		lowerCAmelCase       :     int   				=			{os.path.basename(_lowerCamelCase     ) for p in contents}
		lowerCAmelCase       :     List[Any]   				=			[p for p in contents if p.endswith('ckpt'     )]
		assert len(_lowerCamelCase     ) > 0

		self.assertIn('test_generations.txt'			,       _lowerCamelCase     )
		self.assertIn('test_results.txt'			,       _lowerCamelCase     )

		lowerCAmelCase       :     Tuple   				=			load_json(model.metrics_save_path     )
		lowerCAmelCase       :     Dict   				=			metrics['val'][-1]
		self.assertGreaterEqual(last_step_stats['val_avg_gen_time']			,       0.0_1     )
		self.assertGreaterEqual(1.0			,       last_step_stats['val_avg_gen_time']     )
		self.assertIsInstance(last_step_stats[f"val_avg_{model.val_metric}"]			,       _lowerCamelCase     )
		lowerCAmelCase       :     List[str]   				=			int(args_d['max_epochs'] * (1 / args_d['val_check_interval']) + 1     )
		self.assertEqual(len(metrics['val']     )			,       _lowerCamelCase     )
		self.assertEqual(len(metrics['test']     )			,       1     )
		return model



class 			lowerCAmelCase					(					a_				):



	@parameterized.expand(
	    [T5_TINY, BART_TINY, MBART_TINY, MARIAN_TINY, FSMT_TINY]			,       )
	def 					lowercase  (       self			,       snake_case__     ):
		lowerCAmelCase       :     Any   				=			CHEAP_ARGS.copy()
		lowerCAmelCase       :     Optional[int]   				=			'translation' if model in [MBART_TINY, MARIAN_TINY, FSMT_TINY] else 'summarization'
		lowerCAmelCase       :     str   				=			0.1 if task == 'translation' else 0

		lowerCAmelCase       :     Union[str, Any]   				=			make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()     )
		lowerCAmelCase       :     List[Any]   				=			self.get_auto_remove_tmp_dir()
		args_d.update(
		    data_dir=_lowerCamelCase			,       model_name_or_path=_lowerCamelCase			,       tokenizer_name=_lowerCamelCase			,       train_batch_size=2			,       eval_batch_size=2			,       output_dir=_lowerCamelCase			,       do_predict=_lowerCamelCase			,       task=_lowerCamelCase			,       src_lang='en_XX'			,       tgt_lang='ro_RO'			,       freeze_encoder=_lowerCamelCase			,       freeze_embeds=_lowerCamelCase			,       )
		assert "n_train" in args_d
		lowerCAmelCase       :     Dict   				=			argparse.Namespace(**_lowerCamelCase     )
		lowerCAmelCase       :     Any   				=			main(_lowerCamelCase     )

		lowerCAmelCase       :     Optional[int]   				=			module.model.get_input_embeddings()
		assert not input_embeds.weight.requires_grad
		if model == T5_TINY:
			lowerCAmelCase       :     Optional[Any]   				=			module.model.lm_head
			assert not lm_head.weight.requires_grad
			assert (lm_head.weight == input_embeds.weight).all().item()
		elif model == FSMT_TINY:
			lowerCAmelCase       :     str   				=			module.model.model
			lowerCAmelCase       :     Union[str, Any]   				=			fsmt.decoder.embed_positions
			assert not embed_pos.weight.requires_grad
			assert not fsmt.decoder.embed_tokens.weight.requires_grad
			# check that embeds are not the same
			assert fsmt.decoder.embed_tokens != fsmt.encoder.embed_tokens
		else:
			lowerCAmelCase       :     Dict   				=			module.model.model
			lowerCAmelCase       :     Tuple   				=			bart.decoder.embed_positions
			assert not embed_pos.weight.requires_grad
			assert not bart.shared.weight.requires_grad
			# check that embeds are the same
			assert bart.decoder.embed_tokens == bart.encoder.embed_tokens
			assert bart.decoder.embed_tokens == bart.shared

		lowerCAmelCase       :     int   				=			load_json(module.output_dir / 'text_batch.json'     )
		assert isinstance(_lowerCamelCase			,       _lowerCamelCase     )
		assert len(_lowerCamelCase     ) >= 4



	def 					lowercase  (       self     ):
		lowerCAmelCase       :     int   				=			CHEAP_ARGS.copy()

		lowerCAmelCase       :     int   				=			'summarization'
		lowerCAmelCase       :     List[Any]   				=			make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()     )

		args_d.update(
		    data_dir=_lowerCamelCase			,       tokenizer_name=_lowerCamelCase			,       train_batch_size=2			,       eval_batch_size=2			,       do_predict=_lowerCamelCase			,       task=_lowerCamelCase			,       src_lang='en_XX'			,       tgt_lang='ro_RO'			,       freeze_encoder=_lowerCamelCase			,       freeze_embeds=_lowerCamelCase			,       )

		# test models whose config includes the extra_model_args
		lowerCAmelCase       :     List[str]   				=			BART_TINY
		lowerCAmelCase       :     Union[str, Any]   				=			self.get_auto_remove_tmp_dir()
		lowerCAmelCase       :     List[str]   				=			args_d.copy()
		args_da.update(
		    model_name_or_path=_lowerCamelCase			,       output_dir=_lowerCamelCase			,       )
		lowerCAmelCase       :     Tuple   				=			('encoder_layerdrop', 'decoder_layerdrop', 'dropout', 'attention_dropout')
		for p in extra_model_params:
			lowerCAmelCase       :     Any   				=			0.5
		lowerCAmelCase       :     int   				=			argparse.Namespace(**_lowerCamelCase     )
		lowerCAmelCase       :     Any   				=			main(_lowerCamelCase     )
		for p in extra_model_params:
			assert getattr(model.config			,       _lowerCamelCase     ) == 0.5, f"failed to override the model config for param {p}"

		# test models whose config doesn't include the extra_model_args
		lowerCAmelCase       :     Optional[int]   				=			T5_TINY
		lowerCAmelCase       :     Any   				=			self.get_auto_remove_tmp_dir()
		lowerCAmelCase       :     int   				=			args_d.copy()
		args_da.update(
		    model_name_or_path=_lowerCamelCase			,       output_dir=_lowerCamelCase			,       )
		lowerCAmelCase       :     Optional[int]   				=			'encoder_layerdrop'
		lowerCAmelCase       :     str   				=			0.5
		lowerCAmelCase       :     Optional[Any]   				=			argparse.Namespace(**_lowerCamelCase     )
		with pytest.raises(_lowerCamelCase     ) as excinfo:
			lowerCAmelCase       :     Union[str, Any]   				=			main(_lowerCamelCase     )
		assert str(excinfo.value     ) == f"model config doesn\'t have a `{unsupported_param}` attribute"



	def 					lowercase  (       self     ):
		lowerCAmelCase       :     Union[str, Any]   				=			CHEAP_ARGS.copy()

		lowerCAmelCase       :     Tuple   				=			'summarization'
		lowerCAmelCase       :     Union[str, Any]   				=			make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()     )

		lowerCAmelCase       :     Optional[Any]   				=			BART_TINY
		lowerCAmelCase       :     Dict   				=			self.get_auto_remove_tmp_dir()

		args_d.update(
		    data_dir=_lowerCamelCase			,       model_name_or_path=_lowerCamelCase			,       output_dir=_lowerCamelCase			,       tokenizer_name=_lowerCamelCase			,       train_batch_size=2			,       eval_batch_size=2			,       do_predict=_lowerCamelCase			,       task=_lowerCamelCase			,       src_lang='en_XX'			,       tgt_lang='ro_RO'			,       freeze_encoder=_lowerCamelCase			,       freeze_embeds=_lowerCamelCase			,       )

		# emulate finetune.py
		lowerCAmelCase       :     List[str]   				=			argparse.ArgumentParser()
		lowerCAmelCase       :     Dict   				=			pl.Trainer.add_argparse_args(_lowerCamelCase     )
		lowerCAmelCase       :     List[str]   				=			SummarizationModule.add_model_specific_args(_lowerCamelCase			,       os.getcwd()     )
		lowerCAmelCase       :     List[str]   				=			{'--help': True}

		# --help test
		with pytest.raises(_lowerCamelCase     ) as excinfo:
			with CaptureStdout() as cs:
				lowerCAmelCase       :     Union[str, Any]   				=			parser.parse_args(_lowerCamelCase     )
			assert False, "--help is expected to sys.exit"
		assert excinfo.type == SystemExit
		lowerCAmelCase       :     str   				=			lightning_base.arg_to_scheduler_metavar
		assert expected in cs.out, "--help is expected to list the supported schedulers"

		# --lr_scheduler=non_existing_scheduler test
		lowerCAmelCase       :     Optional[Any]   				=			'non_existing_scheduler'
		lowerCAmelCase       :     Any   				=			{f"--lr_scheduler={unsupported_param}"}
		with pytest.raises(_lowerCamelCase     ) as excinfo:
			with CaptureStderr() as cs:
				lowerCAmelCase       :     Optional[Any]   				=			parser.parse_args(_lowerCamelCase     )
			assert False, "invalid argument is expected to sys.exit"
		assert excinfo.type == SystemExit
		lowerCAmelCase       :     int   				=			f"invalid choice: \'{unsupported_param}\'"
		assert expected in cs.err, f"should have bailed on invalid choice of scheduler {unsupported_param}"

		# --lr_scheduler=existing_scheduler test
		lowerCAmelCase       :     Any   				=			'cosine'
		lowerCAmelCase       :     Tuple   				=			args_d.copy()
		lowerCAmelCase       :     Any   				=			supported_param
		lowerCAmelCase       :     List[Any]   				=			argparse.Namespace(**_lowerCamelCase     )
		lowerCAmelCase       :     Union[str, Any]   				=			main(_lowerCamelCase     )
		assert (
		    getattr(model.hparams			,       'lr_scheduler'     ) == supported_param
		), f"lr_scheduler={supported_param} shouldn\'t fail"








"""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:
        _UpperCamelCase						:    int		=						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)]
                )
                _UpperCamelCase						:    Any		=						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 			a	:
    UpperCAmelCase_					:	Union[str, Any]      	=None
    UpperCAmelCase_					:	Any      	=()
    UpperCAmelCase_					:	Optional[int]      	=()
    UpperCAmelCase_					:	Any      	=True
    UpperCAmelCase_					:	int      	=True
    UpperCAmelCase_					:	Dict      	=True
    UpperCAmelCase_					:	Any      	=False







    def      UpperCamelCase_ (     self				,						_lowerCamelCase				,						_lowerCamelCase				,						_lowerCamelCase=False			):
      lowercase				       =			copy.deepcopy(_lowerCamelCase			)

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

      if return_labels:
        if model_class in get_values(_lowerCamelCase			):
          lowercase				       =			tf.ones(self.model_tester.batch_size				,						dtype=tf.intaa			)
        elif model_class in get_values(_lowerCamelCase			):
          lowercase				       =			tf.zeros(self.model_tester.batch_size				,						dtype=tf.intaa			)
          lowercase				       =			tf.zeros(self.model_tester.batch_size				,						dtype=tf.intaa			)
        elif model_class in [
            *get_values(_lowerCamelCase			),
            *get_values(_lowerCamelCase			),
        ]:
          lowercase				       =			tf.zeros(self.model_tester.batch_size				,						dtype=tf.intaa			)
        elif model_class in get_values(_lowerCamelCase			):
          lowercase				       =			tf.zeros(self.model_tester.batch_size				,						dtype=tf.intaa			)
        elif model_class in [
            *get_values(_lowerCamelCase			),
            *get_values(_lowerCamelCase			),
            *get_values(_lowerCamelCase			),
            *get_values(_lowerCamelCase			),
            *get_values(_lowerCamelCase			),
        ]:
          lowercase				       =			tf.zeros(
              (self.model_tester.batch_size, self.model_tester.seq_length)				,						dtype=tf.intaa			)
      return inputs_dict







    @slow
    def      UpperCamelCase_ (     self			):
      lowercase       ,	lowercase				       =			self.model_tester.prepare_config_and_inputs_for_common()
      for model_class in self.all_model_classes[:2]:
        lowercase				       =			self._prepare_for_class(_lowerCamelCase				,						_lowerCamelCase			)
        lowercase				       =			model_class(_lowerCamelCase			)

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

        lowercase				       =			run_in_graph_mode()
        self.assertIsNotNone(_lowerCamelCase			)







    @slow
    def      UpperCamelCase_ (     self			):
      lowercase       ,	lowercase				       =			self.model_tester.prepare_config_and_inputs_for_common()
      for model_class in self.all_model_classes[:2]:
        lowercase				       =			self._prepare_for_class(_lowerCamelCase				,						_lowerCamelCase			)
        lowercase				       =			model_class(_lowerCamelCase			)

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

        lowercase				       =			run_in_graph_mode()
        self.assertIsNotNone(_lowerCamelCase			)







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

          lowercase				       =			{
              'labels',
              'label',
              'label_ids',
              'start_positions',
              'start_position',
              'end_positions',
              'end_position',
              'next_sentence_label',
          }
          lowercase				       =			possible_label_cols.intersection(set(_lowerCamelCase			)			)
          self.assertGreater(len(_lowerCamelCase			)				,						0				,						msg='No matching label names found!'			)
          lowercase				       =			{key: val for key, val in prepared_for_class.items() if key in label_names}
          lowercase				       =			{key: val for key, val in prepared_for_class.items() if key not in label_names}
          self.assertGreater(len(_lowerCamelCase			)				,						0			)

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

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

          lowercase				       =			history.history['loss'][0]
          self.assertTrue(not isnan(_lowerCamelCase			)			)
          lowercase				       =			history.history['val_loss'][0]
          self.assertTrue(not isnan(_lowerCamelCase			)			)







    @slow
    def      UpperCamelCase_ (     self			):
      lowercase       ,	lowercase				       =			self.model_tester.prepare_config_and_inputs_for_common()
      lowercase				       =			True
      lowercase				       =			True

      if hasattr(_lowerCamelCase				,						'use_cache'			):
        lowercase				       =			True

      lowercase				       =			getattr(self.model_tester				,						'encoder_seq_length'				,						self.model_tester.seq_length			)
      lowercase				       =			getattr(self.model_tester				,						'key_length'				,						_lowerCamelCase			)

      for model_class in self.all_model_classes[:2]:
        lowercase				       =			self._prepare_for_class(_lowerCamelCase				,						_lowerCamelCase			)
        lowercase				       =			model_class(_lowerCamelCase			)
        model.build()
        lowercase				       =			len(model(_lowerCamelCase			)			)

        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:
            lowercase				       =			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(_lowerCamelCase				,						saved_model=_lowerCamelCase			)
          lowercase				       =			os.path.join(_lowerCamelCase				,						'saved_model'				,						'1'			)
          lowercase				       =			tf.keras.models.load_model(_lowerCamelCase			)
          lowercase				       =			model(_lowerCamelCase			)

          if self.is_encoder_decoder:
            lowercase				       =			outputs['encoder_hidden_states']
            lowercase				       =			outputs['encoder_attentions']
          else:
            lowercase				       =			outputs['hidden_states']
            lowercase				       =			outputs['attentions']

          self.assertEqual(len(_lowerCamelCase			)				,						_lowerCamelCase			)

          lowercase				       =			getattr(
              self.model_tester				,						'expected_num_hidden_layers'				,						self.model_tester.num_hidden_layers + 1			)

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

          self.assertEqual(len(_lowerCamelCase			)				,						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      UpperCamelCase_ (     self			):
      tf.keras.mixed_precision.set_global_policy('mixed_float16'			)

      # try/finally block to ensure subsequent tests run in float32
      try:
        lowercase       ,	lowercase				       =			self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes[:2]:
          lowercase				       =			self._prepare_for_class(_lowerCamelCase				,						_lowerCamelCase			)
          lowercase				       =			model_class(_lowerCamelCase			)
          lowercase				       =			model(_lowerCamelCase			)

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







    @slow
    def      UpperCamelCase_ (     self			):
      lowercase       ,	lowercase				       =			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"]
      lowercase				       =			{
          module_member
          for model_class in self.all_model_classes
          for module in (import_module(model_class.__module__			),)
          for module_member_name in dir(_lowerCamelCase			)
          if module_member_name.endswith('MainLayer'			)
          for module_member in (getattr(_lowerCamelCase				,						_lowerCamelCase			),)
          if isinstance(_lowerCamelCase				,						_lowerCamelCase			)
          and tf.keras.layers.Layer in module_member.__bases__
          and getattr(_lowerCamelCase				,						'_keras_serializable'				,						_lowerCamelCase			)
      }

      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
          lowercase				       =			TFSharedEmbeddings(self.model_tester.vocab_size				,						self.model_tester.hidden_size				,						name='shared'			)
          lowercase				       =			False
          lowercase				       =			main_layer_class(_lowerCamelCase				,						embed_tokens=_lowerCamelCase			)
        else:
          lowercase				       =			main_layer_class(_lowerCamelCase			)

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

        if hasattr(self.model_tester				,						'num_labels'			):
          lowercase				       =			self.model_tester.num_labels
        else:
          lowercase				       =			2

        lowercase				       =			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			)

        lowercase				       =			main_layer(_lowerCamelCase			)[0]
        lowercase				       =			tf.keras.layers.Dense(_lowerCamelCase				,						activation='softmax'				,						name='outputs'			)(_lowerCamelCase			)
        lowercase				       =			tf.keras.models.Model(inputs=_lowerCamelCase				,						outputs=[outputs]			)

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

        with tempfile.TemporaryDirectory() as tmpdirname:
          lowercase				       =			os.path.join(_lowerCamelCase				,						'keras_model.h5'			)
          model.save(_lowerCamelCase			)
          if "T5" in main_layer_class.__name__:
            lowercase				       =			tf.keras.models.load_model(
                _lowerCamelCase				,						custom_objects={
                    main_layer_class.__name__: main_layer_class,
                    'TFSharedEmbeddings': TFSharedEmbeddings,
                }				,						)
          else:
            lowercase				       =			tf.keras.models.load_model(
                _lowerCamelCase				,						custom_objects={main_layer_class.__name__: main_layer_class}			)
          assert isinstance(_lowerCamelCase				,						tf.keras.Model			)
          model(_lowerCamelCase			)







    @slow
    def      UpperCamelCase_ (     self			):
      lowercase       ,	lowercase				       =			self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes[:2]:
        lowercase				       =			model_class(_lowerCamelCase			)

        lowercase				       =			copy.deepcopy(_lowerCamelCase			)

        if not self.is_encoder_decoder:
          lowercase				       =			inputs['input_ids']
          del inputs["input_ids"]
        else:
          lowercase				       =			inputs['input_ids']
          lowercase				       =			inputs.get('decoder_input_ids'				,						_lowerCamelCase			)
          del inputs["input_ids"]
          inputs.pop('decoder_input_ids'				,						_lowerCamelCase			)

        if not self.is_encoder_decoder:
          lowercase				       =			model.get_input_embeddings()(_lowerCamelCase			)
        else:
          lowercase				       =			model.get_input_embeddings()(_lowerCamelCase			)
          lowercase				       =			model.get_input_embeddings()(_lowerCamelCase			)

        lowercase				       =			self._prepare_for_class(_lowerCamelCase				,						_lowerCamelCase			)

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

        lowercase				       =			run_in_graph_mode()
        self.assertIsNotNone(_lowerCamelCase			)







    def      UpperCamelCase_ (     self				,						_lowerCamelCase				,						_lowerCamelCase			):
      # special tokens cannot be bad tokens
      lowercase				       =			[]
      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
      lowercase				       =			[]
      while len(_lowerCamelCase			) < num_bad_tokens:
        lowercase				       =			tf.squeeze(ids_tensor((1, 1)				,						self.model_tester.vocab_size			)				,						0			).numpy()[0]
        if token not in special_tokens:
          bad_tokens.append(_lowerCamelCase			)
      return bad_tokens







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







    def      UpperCamelCase_ (     self				,						_lowerCamelCase				,						_lowerCamelCase			):
      # 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(_lowerCamelCase			)				,						len(_lowerCamelCase			)			):
            # if tokens match
            if generated_ids_slice[i - len(_lowerCamelCase			) : i] == bad_word_ids:
              return True
      return False


