




from __future__ import annotations

import unittest

from transformers import AutoTokenizer, GPTJConfig, is_tf_available
from transformers.testing_utils import require_tf, slow, tooslow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin
from ...utils.test_modeling_tf_core import TFCoreModelTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers.models.gptj.modeling_tf_gptj import (
      TFGPTJForCausalLM,
      TFGPTJForQuestionAnswering,
      TFGPTJForSequenceClassification,
      TFGPTJModel,
      shape_list,
  )
class       _UpperCamelCase   :
      def __init__(      self :int			,   lowerCamelCase :List[str]   )  ->      Optional[int]:
       UpperCAmelCase__							=    parent
       UpperCAmelCase__							=    13
       UpperCAmelCase__							=    7
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    99
       UpperCAmelCase__							=    32
       UpperCAmelCase__							=    4
       UpperCAmelCase__							=    2
       UpperCAmelCase__							=    4
       UpperCAmelCase__							=    37
       UpperCAmelCase__							=    "gelu"
       UpperCAmelCase__							=    0.1
       UpperCAmelCase__							=    0.1
       UpperCAmelCase__							=    512
       UpperCAmelCase__							=    16
       UpperCAmelCase__							=    2
       UpperCAmelCase__							=    0.02
       UpperCAmelCase__							=    3
       UpperCAmelCase__							=    4
       UpperCAmelCase__							=    None
       UpperCAmelCase__							=    self.vocab_size - 1
       UpperCAmelCase__							=    self.vocab_size - 1
       UpperCAmelCase__							=    self.vocab_size - 1
      def 						UpperCAmelCase_    (      self :str   )  ->      Union[str, Any]:
       UpperCAmelCase__							=    ids_tensor([self.batch_size, self.seq_length]			,   self.vocab_size   )

       UpperCAmelCase__							=    None
       if self.use_input_mask:
        UpperCAmelCase__							=    random_attention_mask([self.batch_size, self.seq_length]   )

       UpperCAmelCase__							=    None
       if self.use_token_type_ids:
        UpperCAmelCase__							=    ids_tensor([self.batch_size, self.seq_length]			,   self.type_vocab_size   )

       UpperCAmelCase__							=    None
       if self.use_mc_token_ids:
        UpperCAmelCase__							=    ids_tensor([self.batch_size, self.num_choices]			,   self.seq_length   )

       UpperCAmelCase__							=    None
       UpperCAmelCase__							=    None
       UpperCAmelCase__							=    None
       if self.use_labels:
        UpperCAmelCase__							=    ids_tensor([self.batch_size]			,   self.type_sequence_label_size   )
        UpperCAmelCase__							=    ids_tensor([self.batch_size, self.seq_length]			,   self.num_labels   )
        UpperCAmelCase__							=    ids_tensor([self.batch_size]			,   self.num_choices   )

       UpperCAmelCase__							=    GPTJConfig(
           vocab_size=self.vocab_size			,   n_embd=self.hidden_size			,   n_layer=self.num_hidden_layers			,   n_head=self.num_attention_heads			,   intermediate_size=self.intermediate_size			,   hidden_act=self.hidden_act			,   hidden_dropout_prob=self.hidden_dropout_prob			,   attention_probs_dropout_prob=self.attention_probs_dropout_prob			,   n_positions=self.max_position_embeddings			,   type_vocab_size=self.type_vocab_size			,   initializer_range=self.initializer_range			,   bos_token_id=self.bos_token_id			,   eos_token_id=self.eos_token_id			,   pad_token_id=self.pad_token_id			,   rotary_dim=self.rotary_dim			,   return_dict=lowerCamelCase			,   )

       UpperCAmelCase__							=    ids_tensor([self.num_hidden_layers, self.num_attention_heads]			,   2   )

       return (
           config,
           input_ids,
           input_mask,
           head_mask,
           token_type_ids,
           mc_token_ids,
           sequence_labels,
           token_labels,
           choice_labels,
       )
      def 						UpperCAmelCase_    (      self :List[Any]			,   lowerCamelCase :str			,   lowerCamelCase :Tuple			,   lowerCamelCase :Tuple			,   lowerCamelCase :List[str]			,   lowerCamelCase :Any			,   *lowerCamelCase :List[str]   )  ->      Tuple:
       UpperCAmelCase__							=    TFGPTJModel(config=lowerCamelCase   )
       UpperCAmelCase__							=    {
           "input_ids": input_ids,
           "attention_mask": input_mask,
           "token_type_ids": token_type_ids,
       }
       UpperCAmelCase__							=    model(lowerCamelCase   )

       UpperCAmelCase__							=    [input_ids, None, input_mask]  # None is the input for 'past'
       UpperCAmelCase__							=    model(lowerCamelCase   )

       UpperCAmelCase__							=    model(lowerCamelCase   )

       self.parent.assertEqual(result.last_hidden_state.shape			,   (self.batch_size, self.seq_length, self.hidden_size)   )
      def 						UpperCAmelCase_    (      self :str			,   lowerCamelCase :List[Any]			,   lowerCamelCase :str			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :str			,   *lowerCamelCase :int   )  ->      Optional[Any]:
       UpperCAmelCase__							=    TFGPTJModel(config=lowerCamelCase   )

       # first forward pass
       UpperCAmelCase__							=    model(lowerCamelCase			,   token_type_ids=lowerCamelCase			,   use_cache=lowerCamelCase   )
       UpperCAmelCase__							=    model(lowerCamelCase			,   token_type_ids=lowerCamelCase   )
       UpperCAmelCase__							=    model(lowerCamelCase			,   token_type_ids=lowerCamelCase			,   use_cache=lowerCamelCase   )

       self.parent.assertTrue(len(lowerCamelCase   ) == len(lowerCamelCase   )   )
       self.parent.assertTrue(len(lowerCamelCase   ) == len(lowerCamelCase   ) + 1   )

       UpperCAmelCase__			,			UpperCAmelCase__							=    outputs.to_tuple()

       # create hypothetical next token and extent to next_input_ids
       UpperCAmelCase__							=    ids_tensor((self.batch_size, 1)			,   config.vocab_size   )
       UpperCAmelCase__							=    ids_tensor([self.batch_size, 1]			,   self.type_vocab_size   )

       # append to next input_ids and token_type_ids
       UpperCAmelCase__							=    tf.concat([input_ids, next_tokens]			,   axis=-1   )
       UpperCAmelCase__							=    tf.concat([token_type_ids, next_token_types]			,   axis=-1   )

       UpperCAmelCase__							=    model(lowerCamelCase			,   token_type_ids=lowerCamelCase   )["last_hidden_state"]
       UpperCAmelCase__							=    model(lowerCamelCase			,   token_type_ids=lowerCamelCase			,   past_key_values=lowerCamelCase   )[
           "last_hidden_state"
       ]

       # select random slice
       UpperCAmelCase__							=    int(ids_tensor((1,)			,   shape_list(lowerCamelCase   )[-1]   )   )
       UpperCAmelCase__							=    output_from_no_past[:, -1, random_slice_idx]
       UpperCAmelCase__							=    output_from_past[:, 0, random_slice_idx]

       # test that outputs are equal for slice
       tf.debugging.assert_near(lowerCamelCase			,   lowerCamelCase			,   rtol=1e-6   )
      def 						UpperCAmelCase_    (      self :Any			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :Tuple			,   lowerCamelCase :int			,   lowerCamelCase :Dict			,   lowerCamelCase :Any			,   *lowerCamelCase :Union[str, Any]   )  ->      Any:
       UpperCAmelCase__							=    TFGPTJModel(config=lowerCamelCase   )

       # create attention mask
       UpperCAmelCase__							=    self.seq_length // 2
       UpperCAmelCase__							=    tf.ones((self.batch_size, half_seq_length)			,   dtype=tf.intaa   )
       UpperCAmelCase__							=    tf.zeros((self.batch_size, self.seq_length - half_seq_length)			,   dtype=tf.intaa   )
       UpperCAmelCase__							=    tf.concat([attn_mask_begin, attn_mask_end]			,   axis=1   )

       # first forward pass
       UpperCAmelCase__			,			UpperCAmelCase__							=    model(lowerCamelCase			,   attention_mask=lowerCamelCase   ).to_tuple()

       # create hypothetical next token and extent to next_input_ids
       UpperCAmelCase__							=    ids_tensor((self.batch_size, 1)			,   config.vocab_size   )

       # change a random masked slice from input_ids
       UpperCAmelCase__							=    ids_tensor((1,)			,   lowerCamelCase   ).numpy() + 1
       UpperCAmelCase__							=    ids_tensor((self.batch_size, self.seq_length)			,   config.vocab_size   )
       UpperCAmelCase__							=    tf.range(self.seq_length   ) == (self.seq_length - random_seq_idx_to_change)
       UpperCAmelCase__							=    tf.transpose(
           tf.broadcast_to(tf.expand_dims(lowerCamelCase			,   -1   )			,   (self.seq_length, self.batch_size)   )   )
       UpperCAmelCase__							=    tf.where(lowerCamelCase			,   lowerCamelCase			,   lowerCamelCase   )

       # append to next input_ids and attn_mask
       UpperCAmelCase__							=    tf.concat([input_ids, next_tokens]			,   axis=-1   )
       UpperCAmelCase__							=    tf.concat([attn_mask, tf.ones((shape_list(lowerCamelCase   )[0], 1)			,   dtype=tf.intaa   )]			,   axis=1   )

       # get two different outputs
       UpperCAmelCase__							=    model(lowerCamelCase			,   attention_mask=lowerCamelCase   )["last_hidden_state"]
       UpperCAmelCase__							=    model(lowerCamelCase			,   past_key_values=lowerCamelCase			,   attention_mask=lowerCamelCase   )[
           "last_hidden_state"
       ]

       # select random slice
       UpperCAmelCase__							=    int(ids_tensor((1,)			,   shape_list(lowerCamelCase   )[-1]   )   )
       UpperCAmelCase__							=    output_from_no_past[:, -1, random_slice_idx]
       UpperCAmelCase__							=    output_from_past[:, 0, random_slice_idx]

       # test that outputs are equal for slice
       tf.debugging.assert_near(lowerCamelCase			,   lowerCamelCase			,   rtol=1e-12   )
      def 						UpperCAmelCase_    (      self :Dict			,   lowerCamelCase :Any			,   lowerCamelCase :Dict			,   lowerCamelCase :List[str]			,   lowerCamelCase :List[str]			,   lowerCamelCase :Tuple			,   *lowerCamelCase :Optional[int]   )  ->      List[Any]:
       UpperCAmelCase__							=    TFGPTJModel(config=lowerCamelCase   )

       UpperCAmelCase__							=    input_ids[:1, :]
       UpperCAmelCase__							=    input_mask[:1, :]
       UpperCAmelCase__							=    token_type_ids[:1, :]
       UpperCAmelCase__							=    1

       # first forward pass
       UpperCAmelCase__							=    model(lowerCamelCase			,   attention_mask=lowerCamelCase			,   token_type_ids=lowerCamelCase			,   use_cache=lowerCamelCase   )

       UpperCAmelCase__			,			UpperCAmelCase__							=    outputs.to_tuple()

       # create hypothetical next token and extent to next_input_ids
       UpperCAmelCase__							=    ids_tensor((self.batch_size, 3)			,   config.vocab_size   )
       UpperCAmelCase__							=    ids_tensor((self.batch_size, 3)			,   2   )
       UpperCAmelCase__							=    ids_tensor((self.batch_size, 3)			,   self.type_vocab_size   )

       # append to next input_ids and token_type_ids
       UpperCAmelCase__							=    tf.concat([input_ids, next_tokens]			,   axis=-1   )
       UpperCAmelCase__							=    tf.concat([input_mask, next_attn_mask]			,   axis=-1   )
       UpperCAmelCase__							=    tf.concat([token_type_ids, next_token_types]			,   axis=-1   )

       UpperCAmelCase__							=    model(
           lowerCamelCase			,   token_type_ids=lowerCamelCase			,   attention_mask=lowerCamelCase   )["last_hidden_state"]
       UpperCAmelCase__							=    model(
           lowerCamelCase			,   token_type_ids=lowerCamelCase			,   attention_mask=lowerCamelCase			,   past_key_values=lowerCamelCase			,   )["last_hidden_state"]
       self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1]   )

       # select random slice
       UpperCAmelCase__							=    int(ids_tensor((1,)			,   shape_list(lowerCamelCase   )[-1]   )   )
       UpperCAmelCase__							=    output_from_no_past[:, -3:, random_slice_idx]
       UpperCAmelCase__							=    output_from_past[:, :, random_slice_idx]

       # test that outputs are equal for slice
       tf.debugging.assert_near(lowerCamelCase			,   lowerCamelCase			,   rtol=1e-3   )
      def 						UpperCAmelCase_    (      self :Optional[int]			,   lowerCamelCase :Tuple			,   lowerCamelCase :Optional[Any]			,   lowerCamelCase :Tuple			,   lowerCamelCase :int			,   lowerCamelCase :Dict			,   *lowerCamelCase :Optional[Any]   )  ->      List[str]:
       UpperCAmelCase__							=    TFGPTJForCausalLM(config=lowerCamelCase   )
       UpperCAmelCase__							=    {
           "input_ids": input_ids,
           "attention_mask": input_mask,
           "token_type_ids": token_type_ids,
       }
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.seq_length, self.vocab_size)   )







      def 						UpperCAmelCase_    (      self :Union[str, Any]   )  ->      Optional[Any]:
       UpperCAmelCase__							=    self.prepare_config_and_inputs()

       (
           (
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			
       )							=    config_and_inputs

       UpperCAmelCase__							=    {
           "input_ids": input_ids,
           "token_type_ids": token_type_ids,
           "attention_mask": input_mask,
       }
       return config, inputs_dict



@require_tf
class       _UpperCamelCase   (					lowerCAmelCase						,							lowerCAmelCase						,							lowerCAmelCase						,							unittest.TestCase		):
      UpperCAmelCase_          =		(
          (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
          if is_tf_available()
          else ()
      )

      UpperCAmelCase_          =		(TFGPTJForCausalLM,) if is_tf_available() else ()
      UpperCAmelCase_          =		(
          {
              """feature-extraction""": TFGPTJModel,
              """question-answering""": TFGPTJForQuestionAnswering,
              """text-classification""": TFGPTJForSequenceClassification,
              """text-generation""": TFGPTJForCausalLM,
              """zero-shot""": TFGPTJForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      UpperCAmelCase_          =		False
      UpperCAmelCase_          =		False
      UpperCAmelCase_          =		False
      UpperCAmelCase_          =		False
      def 						UpperCAmelCase_    (      self :Union[str, Any]			,   lowerCamelCase :Tuple			,   lowerCamelCase :Tuple			,   lowerCamelCase :List[str]			,   lowerCamelCase :str			,   lowerCamelCase :Union[str, Any]   )  ->      Optional[int]:
       if (
           pipeline_test_casse_name == "QAPipelineTests"
           and tokenizer_name is not None
           and not tokenizer_name.endswith("Fast"   )
       ):
        # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
        # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
        # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
        return True

       return False
      def 						UpperCAmelCase_    (      self :Tuple   )  ->      List[str]:
       UpperCAmelCase__							=    TFGPTJModelTester(self   )
       UpperCAmelCase__							=    ConfigTester(self			,   config_class=lowerCamelCase			,   n_embd=37   )
      def 						UpperCAmelCase_    (      self :str   )  ->      Any:
       self.config_tester.run_common_tests()
      def 						UpperCAmelCase_    (      self :str   )  ->      Tuple:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_gptj_model(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :Any   )  ->      Union[str, Any]:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_gptj_model_past(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :str   )  ->      Optional[Any]:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_gptj_model_attention_mask_past(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :Union[str, Any]   )  ->      Any:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_gptj_model_past_large_inputs(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      Any:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_gptj_lm_head_model(*lowerCamelCase   )
      @slow
      @unittest.skipIf(
          not is_tf_available() or len(tf.config.list_physical_devices("GPU"   )   ) > 0			,   "skip testing on GPU for now to avoid GPU OOM."			,   )
      def 						UpperCAmelCase_    (      self :Any   )  ->      Any:
       UpperCAmelCase__							=    TFGPTJModel.from_pretrained("EleutherAI/gpt-j-6B"			,   from_pt=lowerCamelCase   )
       self.assertIsNotNone(lowerCamelCase   )







      @unittest.skip(reason="Currently, model embeddings are going to undergo a major refactor."   )
      def 						UpperCAmelCase_    (      self :List[str]   )  ->      Dict:
       super().test_resize_token_embeddings()



@require_tf
@tooslow
# Marked as @tooslow due to GPU OOM -- but still useful to run locally. Requires ~39GB of RAM.
class       _UpperCamelCase   (					unittest.TestCase		):
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      Any:
       UpperCAmelCase__							=    TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"			,   from_pt=lowerCamelCase   )
       UpperCAmelCase__							=    tf.convert_to_tensor([[464, 3290]]			,   dtype=tf.intaa   )  # The dog
       # fmt: off
       # The dog is a man's best friend. It is a loyal companion, and it is a friend
       UpperCAmelCase__							=    [464, 3290, 318, 257, 582, 338, 1266, 1545, 13, 632, 318, 257, 9112, 1_5185, 11, 290, 340, 318, 257, 1545]
       # fmt: on
       UpperCAmelCase__							=    model.generate(lowerCamelCase			,   do_sample=lowerCamelCase   )
       self.assertListEqual(output_ids[0].numpy().tolist()			,   lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :Tuple   )  ->      List[str]:
       UpperCAmelCase__							=    AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"			,   revision="float16"   )
       UpperCAmelCase__							=    TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"			,   revision="float16"			,   from_pt=lowerCamelCase   )

       UpperCAmelCase__							=    tokenizer("Today is a nice day and"			,   return_tensors="tf"   )
       # forces the generation to happen on CPU, to avoid GPU-related quirks
       with tf.device(":/CPU:0"   ):
        UpperCAmelCase__							=    model.generate(**lowerCamelCase			,   do_sample=lowerCamelCase			,   seed=[42, 0]   )
       UpperCAmelCase__							=    tokenizer.decode(output_ids[0]			,   skip_special_tokens=lowerCamelCase   )

       UpperCAmelCase__							=    "Today is a nice day and I’m going to go for a walk. I’"
       self.assertEqual(lowerCamelCase			,   lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :List[str]   )  ->      List[Any]:
       UpperCAmelCase__							=    TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"			,   revision="float16"			,   from_pt=lowerCamelCase   )
       UpperCAmelCase__							=    AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"			,   revision="float16"   )

       UpperCAmelCase__							=    "left"

       # Define PAD Token = EOS Token = 50256
       UpperCAmelCase__							=    tokenizer.eos_token
       UpperCAmelCase__							=    model.config.eos_token_id

       # use different length sentences to test batching
       UpperCAmelCase__							=    [
           "Hello, my dog is a little",
           "Today, I",
       ]
       UpperCAmelCase__							=    [
           "Hello, my dog is a little over a year old and has been diagnosed with hip dysplasia",
           "Today, I’m going to be talking about a topic that’",
       ]
       return model, tokenizer, sentences, expected_output_sentences
      def 						UpperCAmelCase_    (      self :Tuple   )  ->      Tuple:
       # Confirms that we get the expected results with left-padded beam search
       UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__							=    self._get_beam_search_test_objects()

       UpperCAmelCase__							=    tokenizer(lowerCamelCase			,   return_tensors="tf"			,   padding=lowerCamelCase   )
       UpperCAmelCase__							=    model.generate(**lowerCamelCase			,   do_sample=lowerCamelCase			,   num_beams=2   )
       UpperCAmelCase__							=    tokenizer.batch_decode(lowerCamelCase			,   skip_special_tokens=lowerCamelCase   )
       self.assertListEqual(lowerCamelCase			,   lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :Optional[int]   )  ->      List[Any]:
       # Confirms that left-padding is working properly
       UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__							=    self._get_beam_search_test_objects()

       UpperCAmelCase__							=    tokenizer(lowerCamelCase			,   return_tensors="tf"			,   padding=lowerCamelCase   )
       UpperCAmelCase__							=    tokenizer(sentences[0]			,   return_tensors="tf"   )
       UpperCAmelCase__							=    model.generate(**lowerCamelCase			,   do_sample=lowerCamelCase			,   num_beams=2   )
       UpperCAmelCase__							=    (
           shape_list(inputs_non_padded["input_ids"]   )[-1]
           - tf.reduce_sum(tf.cast(inputs["attention_mask"][-1]			,   tf.intaa   )   ).numpy()
       )
       UpperCAmelCase__							=    tokenizer(sentences[1]			,   return_tensors="tf"   )
       UpperCAmelCase__							=    model.generate(
           **lowerCamelCase			,   do_sample=lowerCamelCase			,   num_beams=2			,   max_length=model.config.max_length - num_paddings   )
       UpperCAmelCase__							=    tokenizer.decode(output_non_padded[0]			,   skip_special_tokens=lowerCamelCase   )
       UpperCAmelCase__							=    tokenizer.decode(output_padded[0]			,   skip_special_tokens=lowerCamelCase   )
       self.assertListEqual(lowerCamelCase			,   [non_padded_sentence, padded_sentence]   )







      def 						UpperCAmelCase_    (      self :Optional[int]   )  ->      Tuple:
       # Confirms that XLA is working properly
       UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__							=    self._get_beam_search_test_objects()

       UpperCAmelCase__							=    tokenizer(lowerCamelCase			,   return_tensors="tf"			,   padding=lowerCamelCase   )
       UpperCAmelCase__							=    tf.function(model.generate			,   jit_compile=lowerCamelCase   )
       UpperCAmelCase__							=    xla_generate(**lowerCamelCase			,   do_sample=lowerCamelCase			,   num_beams=2   )
       UpperCAmelCase__							=    tokenizer.batch_decode(lowerCamelCase			,   skip_special_tokens=lowerCamelCase   )
       self.assertListEqual(lowerCamelCase			,   lowerCamelCase   )











from __future__ import annotations

import unittest

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

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers import TF_MODEL_FOR_PRETRAINING_MAPPING
  from transformers.models.albert.modeling_tf_albert import (
      TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFAlbertForMaskedLM,
      TFAlbertForMultipleChoice,
      TFAlbertForPreTraining,
      TFAlbertForQuestionAnswering,
      TFAlbertForSequenceClassification,
      TFAlbertForTokenClassification,
      TFAlbertModel,
  )
class       _UpperCamelCase   :
      def __init__(      self :Optional[int]			,   lowerCamelCase :Tuple			,   lowerCamelCase :Union[str, Any]=13			,   lowerCamelCase :List[str]=7			,   lowerCamelCase :List[Any]=True			,   lowerCamelCase :Any=True			,   lowerCamelCase :Any=True			,   lowerCamelCase :Optional[Any]=True			,   lowerCamelCase :Optional[int]=99			,   lowerCamelCase :Tuple=16			,   lowerCamelCase :Optional[Any]=32			,   lowerCamelCase :Optional[Any]=2			,   lowerCamelCase :List[Any]=4			,   lowerCamelCase :Dict=37			,   lowerCamelCase :Optional[Any]="gelu"			,   lowerCamelCase :Dict=0.1			,   lowerCamelCase :Optional[Any]=0.1			,   lowerCamelCase :Tuple=512			,   lowerCamelCase :Optional[int]=16			,   lowerCamelCase :Any=2			,   lowerCamelCase :Tuple=0.02			,   lowerCamelCase :Union[str, Any]=3			,   lowerCamelCase :str=4			,   lowerCamelCase :Dict=None			,   )  ->      List[Any]:
       UpperCAmelCase__							=    parent
       UpperCAmelCase__							=    13
       UpperCAmelCase__							=    7
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    True
       UpperCAmelCase__							=    99
       UpperCAmelCase__							=    16
       UpperCAmelCase__							=    32
       UpperCAmelCase__							=    2
       UpperCAmelCase__							=    4
       UpperCAmelCase__							=    37
       UpperCAmelCase__							=    "gelu"
       UpperCAmelCase__							=    0.1
       UpperCAmelCase__							=    0.1
       UpperCAmelCase__							=    512
       UpperCAmelCase__							=    16
       UpperCAmelCase__							=    2
       UpperCAmelCase__							=    0.02
       UpperCAmelCase__							=    3
       UpperCAmelCase__							=    4
       UpperCAmelCase__							=    None
      def 						UpperCAmelCase_    (      self :Optional[int]   )  ->      Dict:
       UpperCAmelCase__							=    ids_tensor([self.batch_size, self.seq_length]			,   self.vocab_size   )

       UpperCAmelCase__							=    None
       if self.use_input_mask:
        UpperCAmelCase__							=    random_attention_mask([self.batch_size, self.seq_length]   )

       UpperCAmelCase__							=    None
       if self.use_token_type_ids:
        UpperCAmelCase__							=    ids_tensor([self.batch_size, self.seq_length]			,   self.type_vocab_size   )

       UpperCAmelCase__							=    None
       UpperCAmelCase__							=    None
       UpperCAmelCase__							=    None
       if self.use_labels:
        UpperCAmelCase__							=    ids_tensor([self.batch_size]			,   self.type_sequence_label_size   )
        UpperCAmelCase__							=    ids_tensor([self.batch_size, self.seq_length]			,   self.num_labels   )
        UpperCAmelCase__							=    ids_tensor([self.batch_size]			,   self.num_choices   )

       UpperCAmelCase__							=    AlbertConfig(
           vocab_size=self.vocab_size			,   hidden_size=self.hidden_size			,   embedding_size=self.embedding_size			,   num_hidden_layers=self.num_hidden_layers			,   num_attention_heads=self.num_attention_heads			,   intermediate_size=self.intermediate_size			,   hidden_act=self.hidden_act			,   hidden_dropout_prob=self.hidden_dropout_prob			,   attention_probs_dropout_prob=self.attention_probs_dropout_prob			,   max_position_embeddings=self.max_position_embeddings			,   type_vocab_size=self.type_vocab_size			,   initializer_range=self.initializer_range			,   )

       return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
      def 						UpperCAmelCase_    (      self :Optional[int]			,   lowerCamelCase :str			,   lowerCamelCase :Tuple			,   lowerCamelCase :Tuple			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :int			,   lowerCamelCase :str   )  ->      Optional[Any]:
       UpperCAmelCase__							=    TFAlbertModel(config=lowerCamelCase   )
       # inputs = {'input_ids': input_ids,
       #           'attention_mask': input_mask,
       #           'token_type_ids': token_type_ids}
       # sequence_output, pooled_output = model(**inputs)
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(lowerCamelCase   )

       UpperCAmelCase__							=    [input_ids, input_mask]
       UpperCAmelCase__							=    model(lowerCamelCase   )

       UpperCAmelCase__							=    model(lowerCamelCase   )

       self.parent.assertEqual(result.last_hidden_state.shape			,   (self.batch_size, self.seq_length, self.hidden_size)   )
       self.parent.assertEqual(result.pooler_output.shape			,   (self.batch_size, self.hidden_size)   )
      def 						UpperCAmelCase_    (      self :Optional[Any]			,   lowerCamelCase :Any			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Any			,   lowerCamelCase :List[str]			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Tuple   )  ->      Optional[int]:
       UpperCAmelCase__							=    self.num_labels
       UpperCAmelCase__							=    TFAlbertForPreTraining(config=lowerCamelCase   )
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertEqual(result.prediction_logits.shape			,   (self.batch_size, self.seq_length, self.vocab_size)   )
       self.parent.assertEqual(result.sop_logits.shape			,   (self.batch_size, self.num_labels)   )
      def 						UpperCAmelCase_    (      self :int			,   lowerCamelCase :Dict			,   lowerCamelCase :List[Any]			,   lowerCamelCase :List[str]			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :Tuple			,   lowerCamelCase :Any   )  ->      Optional[Any]:
       UpperCAmelCase__							=    TFAlbertForMaskedLM(config=lowerCamelCase   )
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.seq_length, self.vocab_size)   )
      def 						UpperCAmelCase_    (      self :Tuple			,   lowerCamelCase :str			,   lowerCamelCase :Dict			,   lowerCamelCase :Dict			,   lowerCamelCase :Optional[Any]			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Optional[Any]   )  ->      str:
       UpperCAmelCase__							=    self.num_labels
       UpperCAmelCase__							=    TFAlbertForSequenceClassification(config=lowerCamelCase   )
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.num_labels)   )
      def 						UpperCAmelCase_    (      self :Dict			,   lowerCamelCase :Optional[Any]			,   lowerCamelCase :str			,   lowerCamelCase :Dict			,   lowerCamelCase :List[Any]			,   lowerCamelCase :List[str]			,   lowerCamelCase :int			,   lowerCamelCase :Optional[Any]   )  ->      Dict:
       UpperCAmelCase__							=    TFAlbertForQuestionAnswering(config=lowerCamelCase   )
       UpperCAmelCase__							=    {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertEqual(result.start_logits.shape			,   (self.batch_size, self.seq_length)   )
       self.parent.assertEqual(result.end_logits.shape			,   (self.batch_size, self.seq_length)   )
      def 						UpperCAmelCase_    (      self :Optional[int]			,   lowerCamelCase :Optional[Any]			,   lowerCamelCase :Dict			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Tuple			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Union[str, Any]   )  ->      int:
       UpperCAmelCase__							=    self.num_choices
       UpperCAmelCase__							=    TFAlbertForMultipleChoice(config=lowerCamelCase   )
       UpperCAmelCase__							=    tf.tile(tf.expand_dims(lowerCamelCase			,   1   )			,   (1, self.num_choices, 1)   )
       UpperCAmelCase__							=    tf.tile(tf.expand_dims(lowerCamelCase			,   1   )			,   (1, self.num_choices, 1)   )
       UpperCAmelCase__							=    tf.tile(tf.expand_dims(lowerCamelCase			,   1   )			,   (1, self.num_choices, 1)   )
       UpperCAmelCase__							=    {
           "input_ids": multiple_choice_inputs_ids,
           "attention_mask": multiple_choice_input_mask,
           "token_type_ids": multiple_choice_token_type_ids,
       }
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertListEqual(list(result["logits"].shape   )			,   [self.batch_size, self.num_choices]   )
      def 						UpperCAmelCase_    (      self :int			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :List[str]			,   lowerCamelCase :int			,   lowerCamelCase :int			,   lowerCamelCase :Union[str, Any]			,   lowerCamelCase :Dict			,   lowerCamelCase :Dict   )  ->      List[str]:
       UpperCAmelCase__							=    self.num_labels
       UpperCAmelCase__							=    TFAlbertForTokenClassification(config=lowerCamelCase   )
       UpperCAmelCase__							=    {
           "input_ids": input_ids,
           "attention_mask": input_mask,
           "token_type_ids": token_type_ids,
       }
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertListEqual(list(result["logits"].shape   )			,   [self.batch_size, self.seq_length, self.num_labels]   )







      def 						UpperCAmelCase_    (      self :str   )  ->      Any:
       UpperCAmelCase__							=    self.prepare_config_and_inputs()
       (
           (
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			
       )							=    config_and_inputs
       UpperCAmelCase__							=    {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
       return config, inputs_dict



@require_tf
class       _UpperCamelCase   (					lowerCAmelCase						,							lowerCAmelCase						,							unittest.TestCase		):
      UpperCAmelCase_          =		(
          (
              TFAlbertModel,
              TFAlbertForPreTraining,
              TFAlbertForMaskedLM,
              TFAlbertForSequenceClassification,
              TFAlbertForQuestionAnswering,
              TFAlbertForTokenClassification,
              TFAlbertForMultipleChoice,
          )
          if is_tf_available()
          else ()
      )
      UpperCAmelCase_          =		(
          {
              """feature-extraction""": TFAlbertModel,
              """fill-mask""": TFAlbertForMaskedLM,
              """question-answering""": TFAlbertForQuestionAnswering,
              """text-classification""": TFAlbertForSequenceClassification,
              """token-classification""": TFAlbertForTokenClassification,
              """zero-shot""": TFAlbertForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      UpperCAmelCase_          =		False
      UpperCAmelCase_          =		False
      def 						UpperCAmelCase_    (      self :str			,   lowerCamelCase :Any			,   lowerCamelCase :Optional[int]			,   lowerCamelCase :Optional[int]=False   )  ->      str:
       UpperCAmelCase__							=    super()._prepare_for_class(lowerCamelCase			,   lowerCamelCase			,   return_labels=lowerCamelCase   )

       if return_labels:
        if model_class in get_values(lowerCamelCase   ):
         UpperCAmelCase__							=    tf.zeros(self.model_tester.batch_size			,   dtype=tf.intaa   )

       return inputs_dict
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      List[Any]:
       UpperCAmelCase__							=    TFAlbertModelTester(self   )
       UpperCAmelCase__							=    ConfigTester(self			,   config_class=lowerCamelCase			,   hidden_size=37   )
      def 						UpperCAmelCase_    (      self :Dict   )  ->      Union[str, Any]:
       self.config_tester.run_common_tests()
      def 						UpperCAmelCase_    (      self :Optional[Any]   )  ->      Tuple:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_model(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      Dict:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_pretraining(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      str:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_masked_lm(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :Dict   )  ->      int:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_multiple_choice(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :Union[str, Any]   )  ->      Dict:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_sequence_classification(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :str   )  ->      int:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_albert_for_question_answering(*lowerCamelCase   )







      @slow
      def 						UpperCAmelCase_    (      self :int   )  ->      str:
       for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCAmelCase__							=    TFAlbertModel.from_pretrained(lowerCamelCase   )
        self.assertIsNotNone(lowerCamelCase   )



@require_tf
class       _UpperCamelCase   (					unittest.TestCase		):
      @slow
      def 						UpperCAmelCase_    (      self :List[str]   )  ->      List[str]:
       UpperCAmelCase__							=    TFAlbertForPreTraining.from_pretrained("albert-base-v2"   )
       UpperCAmelCase__							=    tf.constant([[0, 1, 2, 3, 4, 5]]   )
       UpperCAmelCase__							=    model(lowerCamelCase   )[0]

       UpperCAmelCase__							=    [1, 6, 3_0000]
       self.assertEqual(output.shape			,   lowerCamelCase   )

       UpperCAmelCase__							=    tf.constant(
           [
               [
                   [4.59_56_68, 0.74_46_27_54, -1.81_81_47],
                   [4.5_95_43_47, 0.7_45_41_84, -1.8_18_82_58],
                   [4.5_95_49_05, 0.7_44_82_35, -1.8_18_23_16],
               ]
           ]   )
       tf.debugging.assert_near(output[:, :3, :3]			,   lowerCamelCase			,   atol=1e-4   )





