






"""simple docstring"""



import argparse
import inspect
import logging
from typing import Tuple

import numpy as np
import torch

from transformers import (
    AutoTokenizer,
    BloomForCausalLM,
    BloomTokenizerFast,
    CTRLLMHeadModel,
    CTRLTokenizer,
    GenerationMixin,
    GPTaLMHeadModel,
    GPTaTokenizer,
    GPTJForCausalLM,
    LlamaForCausalLM,
    LlamaTokenizer,
    OpenAIGPTLMHeadModel,
    OpenAIGPTTokenizer,
    OPTForCausalLM,
    TransfoXLLMHeadModel,
    TransfoXLTokenizer,
    XLMTokenizer,
    XLMWithLMHeadModel,
    XLNetLMHeadModel,
    XLNetTokenizer,
)
from transformers.modeling_outputs import CausalLMOutputWithPast


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
    datefmt='%m/%d/%Y %H:%M:%S',
    level=logging.INFO,
)
_UpperCamelCase						:    Optional[int]		=						logging.getLogger(__name__)

_UpperCamelCase						:    Any		=						int(1_0_0_0_0)  # Hardcoded max length to avoid infinite loop

_UpperCamelCase						:    int		=						{
    'gpt2': (GPTaLMHeadModel, GPTaTokenizer),
    'ctrl': (CTRLLMHeadModel, CTRLTokenizer),
    'openai-gpt': (OpenAIGPTLMHeadModel, OpenAIGPTTokenizer),
    'xlnet': (XLNetLMHeadModel, XLNetTokenizer),
    'transfo-xl': (TransfoXLLMHeadModel, TransfoXLTokenizer),
    'xlm': (XLMWithLMHeadModel, XLMTokenizer),
    'gptj': (GPTJForCausalLM, AutoTokenizer),
    'bloom': (BloomForCausalLM, BloomTokenizerFast),
    'llama': (LlamaForCausalLM, LlamaTokenizer),
    'opt': (OPTForCausalLM, GPTaTokenizer),
}

# Padding text to help Transformer-XL and XLNet with short prompts as proposed by Aman Rusia
# in https://github.com/rusiaaman/XLNet-gen#methodology
# and https://medium.com/@amanrusia/xlnet-speaks-comparison-to-gpt-2-ea1a4e9ba39e
_UpperCamelCase						:    Optional[Any]		=						'In 1991, the remains of Russian Tsar Nicholas II and his family\n(except for Alexei and Maria) are discovered.\nThe voice of Nicholas\'s young son, Tsarevich Alexei Nikolaevich, narrates the\nremainder of the story. 1883 Western Siberia,\na young Grigori Rasputin is asked by his father and a group of men to perform magic.\nRasputin has a vision and denounces one of the men as a horse thief. Although his\nfather initially slaps him for making such an accusation, Rasputin watches as the\nman is chased outside and beaten. Twenty years later, Rasputin sees a vision of\nthe Virgin Mary, prompting him to become a priest. Rasputin quickly becomes famous,\nwith people, even a bishop, begging for his blessing. <eod> </s> <eos>'
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    List[str]       ):





  '''simple docstring'''





  np.random.seed(args.seed       )
  torch.manual_seed(args.seed       )
  if args.n_gpu > 0:
    torch.cuda.manual_seed_all(args.seed       )
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Any						,   __snake_case     :    Any						,   __snake_case     :    Union[str, Any]						,   __snake_case     :    List[Any]       ):





  '''simple docstring'''





  if args.temperature > 0.7:
    logger.info('CTRL typically works better with lower temperatures (and lower top_k).'       )

  lowercase				       =			tokenizer.encode(__snake_case						,   add_special_tokens=__snake_case       )
  if not any(encoded_prompt[0] == x for x in tokenizer.control_codes.values()       ):
    logger.info('WARNING! You are not starting your generation from a control code so you won\'t get good results'       )
  return prompt_text
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Tuple						,   __snake_case     :    List[Any]						,   __snake_case     :    Any						,   __snake_case     :    str       ):





  '''simple docstring'''





  lowercase				       =			hasattr(model.config						,   'use_lang_emb'       ) and model.config.use_lang_emb
  if hasattr(model.config						,   'lang2id'       ) and use_lang_emb:
    lowercase				       =			model.config.langaid.keys()
    if args.xlm_language in available_languages:
      lowercase				       =			args.xlm_language
    else:
      lowercase				       =			None
      while language not in available_languages:
        lowercase				       =			input('Using XLM. Select language in ' + str(list(__snake_case       )       ) + ' >>> '       )

    lowercase				       =			model.config.langaid[language]
    # kwargs["language"] = tokenizer.lang2id[language]

  # TODO fix mask_token_id setup when configurations will be synchronized between models and tokenizers
  # XLM masked-language modeling (MLM) models need masked token
  # is_xlm_mlm = "mlm" in args.model_name_or_path
  # if is_xlm_mlm:
  #     kwargs["mask_token_id"] = tokenizer.mask_token_id

  return prompt_text
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Optional[int]						,   __snake_case     :    int						,   __snake_case     :    List[Any]						,   __snake_case     :    int       ):





  '''simple docstring'''





  lowercase				       =			args.prefix if args.prefix else args.padding_text if args.padding_text else PREFIX
  lowercase				       =			prefix + prompt_text
  return prompt_text
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    int						,   __snake_case     :    Optional[int]						,   __snake_case     :    Dict						,   __snake_case     :    Any       ):





  '''simple docstring'''





  lowercase				       =			args.prefix if args.prefix else args.padding_text if args.padding_text else PREFIX
  lowercase				       =			prefix + prompt_text
  return prompt_text


_UpperCamelCase						:    str		=						{
    'ctrl': prepare_ctrl_input,
    'xlm': prepare_xlm_input,
    'xlnet': prepare_xlnet_input,
    'transfo-xl': prepare_transfoxl_input,
}
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Optional[int]						,   __snake_case     :    str       ):





  '''simple docstring'''





  if length < 0 and max_sequence_length > 0:
    lowercase				       =			max_sequence_length
  elif 0 < max_sequence_length < length:
    lowercase				       =			max_sequence_length  # No generation bigger than model size
  elif length < 0:
    lowercase				       =			MAX_LENGTH  # avoid infinite loop
  return length
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Optional[Any]       ):





  '''simple docstring'''





  lowercase				       =			None
  if hasattr(__snake_case						,   'hidden_size'       ):
    lowercase				       =			model_config.hidden_size
  elif hasattr(__snake_case						,   'n_embed'       ):
    lowercase				       =			model_config.n_embed
  elif hasattr(__snake_case						,   'n_embd'       ):
    lowercase				       =			model_config.n_embd

  lowercase				       =			None
  if hasattr(__snake_case						,   'num_attention_heads'       ):
    lowercase				       =			model_config.num_attention_heads
  elif hasattr(__snake_case						,   'n_head'       ):
    lowercase				       =			model_config.n_head

  if embedding_size is None or num_head is None or num_head == 0:
    raise ValueError('Check the model config'       )

  lowercase				       =			int(embedding_size / num_head       )
  if hasattr(__snake_case						,   'n_layer'       ):
    lowercase				       =			model_config.n_layer
  elif hasattr(__snake_case						,   'num_hidden_layers'       ):
    lowercase				       =			model_config.num_hidden_layers
  else:
    raise ValueError('Number of hidden layers couldn\'t be determined from the model config'       )

  return num_layer, num_head, num_embedding_size_per_head
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    List[Any]						,   __snake_case     :    int						,   __snake_case     :    str       ):





  '''simple docstring'''





  lowercase       ,	lowercase       ,	lowercase				       =			sparse_model_config(model.config       )
  if model.config.model_type == "bloom":
    lowercase				       =			tuple(
        (
            torch.empty(int(num_attention_heads * batch_size       )						,   __snake_case						,   __snake_case       )
            .to(model.dtype       )
            .to(model.device       ),
            torch.empty(int(num_attention_heads * batch_size       )						,   __snake_case						,   __snake_case       )
            .to(model.dtype       )
            .to(model.device       ),
        )
        for _ in range(__snake_case       )       )
  else:
    lowercase				       =			tuple(
        (
            torch.empty(__snake_case						,   __snake_case						,   __snake_case						,   __snake_case       )
            .to(model.dtype       )
            .to(model.device       ),
            torch.empty(__snake_case						,   __snake_case						,   __snake_case						,   __snake_case       )
            .to(model.dtype       )
            .to(model.device       ),
        )
        for _ in range(__snake_case       )       )
  return past_key_values
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Union[str, Any]						,   __snake_case     :    int						,   __snake_case     :    int       ):





  '''simple docstring'''





  lowercase				       =			len(__snake_case       )
  lowercase				       =			tokenizer.batch_encode_plus(__snake_case						,   return_tensors='pt'       )
  lowercase				       =			dummy_input.to(model.device       )
  if model.config.use_cache:
    lowercase				       =			generate_past_key_values(__snake_case						,   __snake_case						,   1       )
  lowercase				       =			torch.cat(
      [
          torch.zeros(dummy_input['attention_mask'].shape[0]						,   1       )
          .to(dummy_input['attention_mask'].dtype       )
          .to(model.device       ),
          dummy_input['attention_mask'],
      ]						,   -1						,   )
  return dummy_input

class 			a	(  a_      ):
    UpperCAmelCase_					:	str      	=("_optimized", "_default")







    def __init__(     self				,						_lowerCamelCase				,						_lowerCamelCase			):
      lowercase				       =			optimized
      lowercase				       =			default







    def __call__(     self				,						*_lowerCamelCase				,						**_lowerCamelCase			):
      if kwargs["past_key_values"] is None and self._default.config.use_cache:
        lowercase				       =			generate_past_key_values(self._default				,						kwargs['input_ids'].shape[0]				,						0			)
      kwargs.pop('position_ids'				,						_lowerCamelCase			)
      for k in list(kwargs.keys()			):
        if kwargs[k] is None or isinstance(kwargs[k]				,						_lowerCamelCase			):
          kwargs.pop(_lowerCamelCase			)
      lowercase				       =			self._optimized(**_lowerCamelCase			)
      lowercase				       =			outputs[0]
      lowercase				       =			outputs[1]
      lowercase				       =			CausalLMOutputWithPast(
          loss=_lowerCamelCase				,						logits=_lowerCamelCase				,						past_key_values=_lowerCamelCase				,						hidden_states=_lowerCamelCase				,						attentions=_lowerCamelCase				,						)
      return fixed_output







    def __getattr__(     self				,						_lowerCamelCase			):
      return getattr(self._default				,						_lowerCamelCase			)







    def      UpperCamelCase_ (     self				,						_lowerCamelCase				,						_lowerCamelCase=None				,						_lowerCamelCase=None				,						_lowerCamelCase=None				,						**_lowerCamelCase			):
      return self._default.prepare_inputs_for_generation(
          _lowerCamelCase				,						past_key_values=_lowerCamelCase				,						inputs_embeds=_lowerCamelCase				,						use_cache=_lowerCamelCase				,						**_lowerCamelCase			)







    def      UpperCamelCase_ (     self				,						_lowerCamelCase				,						_lowerCamelCase			):
      return self._default._reorder_cache(_lowerCamelCase				,						_lowerCamelCase			)


def 		_SCREAMING_SNAKE_CASE					(       ):





  '''simple docstring'''





  lowercase				       =			argparse.ArgumentParser()
  parser.add_argument(
      '--model_type'						,   default=__snake_case						,   type=__snake_case						,   required=__snake_case						,   help='Model type selected in the list: ' + ', '.join(MODEL_CLASSES.keys()       )						,   )
  parser.add_argument(
      '--model_name_or_path'						,   default=__snake_case						,   type=__snake_case						,   required=__snake_case						,   help='Path to pre-trained model or shortcut name selected in the list: ' + ', '.join(MODEL_CLASSES.keys()       )						,   )

  parser.add_argument('--prompt'						,   type=__snake_case						,   default=''       )
  parser.add_argument('--length'						,   type=__snake_case						,   default=20       )
  parser.add_argument('--stop_token'						,   type=__snake_case						,   default=__snake_case						,   help='Token at which text generation is stopped'       )

  parser.add_argument(
      '--temperature'						,   type=__snake_case						,   default=1.0						,   help='temperature of 1.0 has no effect, lower tend toward greedy sampling'						,   )
  parser.add_argument(
      '--repetition_penalty'						,   type=__snake_case						,   default=1.0						,   help='primarily useful for CTRL model; in that case, use 1.2'       )
  parser.add_argument('--k'						,   type=__snake_case						,   default=0       )
  parser.add_argument('--p'						,   type=__snake_case						,   default=0.9       )

  parser.add_argument('--prefix'						,   type=__snake_case						,   default=''						,   help='Text added prior to input.'       )
  parser.add_argument('--padding_text'						,   type=__snake_case						,   default=''						,   help='Deprecated, the use of `--prefix` is preferred.'       )
  parser.add_argument('--xlm_language'						,   type=__snake_case						,   default=''						,   help='Optional language when used with the XLM model.'       )

  parser.add_argument('--seed'						,   type=__snake_case						,   default=42						,   help='random seed for initialization'       )
  parser.add_argument('--no_cuda'						,   action='store_true'						,   help='Avoid using CUDA when available'       )
  parser.add_argument('--num_return_sequences'						,   type=__snake_case						,   default=1						,   help='The number of samples to generate.'       )
  parser.add_argument(
      '--fp16'						,   action='store_true'						,   help='Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit'						,   )
  parser.add_argument('--jit'						,   action='store_true'						,   help='Whether or not to use jit trace to accelerate inference'       )
  lowercase				       =			parser.parse_args()

  lowercase				       =			torch.device('cuda' if torch.cuda.is_available() and not args.no_cuda else 'cpu'       )
  lowercase				       =			0 if args.no_cuda else torch.cuda.device_count()

  logger.warning(f'device: {args.device}, n_gpu: {args.n_gpu}, 16-bits training: {args.fpaa}'       )

  set_seed(__snake_case       )

  # Initialize the model and tokenizer
  try:
    lowercase				       =			args.model_type.lower()
    lowercase       ,	lowercase				       =			MODEL_CLASSES[args.model_type]
  except KeyError:
    raise KeyError('the model {} you specified is not supported. You are welcome to add it and open a PR :)'       )

  lowercase				       =			tokenizer_class.from_pretrained(args.model_name_or_path       )
  if tokenizer.pad_token is None:
    lowercase				       =			tokenizer.eos_token
  lowercase				       =			model_class.from_pretrained(args.model_name_or_path       )
  model.to(args.device       )

  if args.fpaa:
    model.half()
  lowercase				       =			getattr(model.config						,   'max_position_embeddings'						,   0       )
  lowercase				       =			adjust_length_to_model(args.length						,   max_sequence_length=__snake_case       )
  logger.info(__snake_case       )

  lowercase				       =			args.prompt if args.prompt else input('Model prompt >>> '       )

  # Different models need different input formatting and/or extra arguments
  lowercase				       =			args.model_type in PREPROCESSING_FUNCTIONS.keys()
  if requires_preprocessing:
    lowercase				       =			PREPROCESSING_FUNCTIONS.get(args.model_type       )
    lowercase				       =			prepare_input(__snake_case						,   __snake_case						,   __snake_case						,   __snake_case       )

    if model.__class__.__name__ in ["TransfoXLLMHeadModel"]:
      lowercase				       =			{'add_space_before_punct_symbol': True}
    else:
      lowercase				       =			{}

    lowercase				       =			tokenizer.encode(
        __snake_case						,   add_special_tokens=__snake_case						,   return_tensors='pt'						,   **__snake_case       )
  else:
    lowercase				       =			args.prefix if args.prefix else args.padding_text
    lowercase				       =			tokenizer.encode(prefix + prompt_text						,   add_special_tokens=__snake_case						,   return_tensors='pt'       )
  lowercase				       =			encoded_prompt.to(args.device       )

  if encoded_prompt.size()[-1] == 0:
    lowercase				       =			None
  else:
    lowercase				       =			encoded_prompt

  if args.jit:
    lowercase				       =			['enable jit']
    lowercase				       =			prepare_jit_inputs(__snake_case						,   __snake_case						,   __snake_case       )
    torch._C._jit_set_texpr_fuser_enabled(__snake_case       )
    lowercase				       =			False
    if hasattr(__snake_case						,   'forward'       ):
      lowercase				       =			inspect.signature(model.forward       )
    else:
      lowercase				       =			inspect.signature(model.__call__       )
    lowercase				       =			tuple(jit_inputs[key] for key in sig.parameters if jit_inputs.get(__snake_case						,   __snake_case       ) is not None       )
    lowercase				       =			torch.jit.trace(__snake_case						,   __snake_case						,   strict=__snake_case       )
    lowercase				       =			torch.jit.freeze(traced_model.eval()       )
    traced_model(*__snake_case       )
    traced_model(*__snake_case       )

    lowercase				       =			_ModelFallbackWrapper(__snake_case						,   __snake_case       )

  lowercase				       =			model.generate(
      input_ids=__snake_case						,   max_length=args.length + len(encoded_prompt[0]       )						,   temperature=args.temperature						,   top_k=args.k						,   top_p=args.p						,   repetition_penalty=args.repetition_penalty						,   do_sample=__snake_case						,   num_return_sequences=args.num_return_sequences						,   )

  # Remove the batch dimension when returning multiple sequences
  if len(output_sequences.shape       ) > 2:
    output_sequences.squeeze_()

  lowercase				       =			[]

  for generated_sequence_idx, generated_sequence in enumerate(__snake_case       ):
    print(f'=== GENERATED SEQUENCE {generated_sequence_idx + 1} ==='       )
    lowercase				       =			generated_sequence.tolist()

    # Decode text
    lowercase				       =			tokenizer.decode(__snake_case						,   clean_up_tokenization_spaces=__snake_case       )

    # Remove all text after the stop token
    lowercase				       =			text[: text.find(args.stop_token       ) if args.stop_token else None]

    # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
    lowercase				       =			(
        prompt_text + text[len(tokenizer.decode(encoded_prompt[0]						,   clean_up_tokenization_spaces=__snake_case       )       ) :]
    )

    generated_sequences.append(__snake_case       )
    print(__snake_case       )

  return generated_sequences


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


