

'''simple docstring'''






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




    def __init__( self    ,					lowercase__      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      Tuple					=	parent
          SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	13
          SCREAMING_SNAKE_CASE_				:      List[str]					=	7
          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	True
          SCREAMING_SNAKE_CASE_				:      List[str]					=	True
          SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	True
          SCREAMING_SNAKE_CASE_				:      List[Any]					=	True
          SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	True
          SCREAMING_SNAKE_CASE_				:      Tuple					=	99
          SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	32
          SCREAMING_SNAKE_CASE_				:      int					=	4
          SCREAMING_SNAKE_CASE_				:      List[Any]					=	2
          SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	4
          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	37
          SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	"gelu"
          SCREAMING_SNAKE_CASE_				:      int					=	0.1
          SCREAMING_SNAKE_CASE_				:      Any					=	0.1
          SCREAMING_SNAKE_CASE_				:      List[Any]					=	512
          SCREAMING_SNAKE_CASE_				:      int					=	16
          SCREAMING_SNAKE_CASE_				:      List[str]					=	2
          SCREAMING_SNAKE_CASE_				:      Any					=	0.02
          SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	3
          SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	4
          SCREAMING_SNAKE_CASE_				:      List[Any]					=	None
          SCREAMING_SNAKE_CASE_				:      Tuple					=	self.vocab_size - 1
          SCREAMING_SNAKE_CASE_				:      str					=	self.vocab_size - 1
          SCREAMING_SNAKE_CASE_				:      str					=	self.vocab_size - 1




    def   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      List[Any]					=	ids_tensor([self.batch_size, self.seq_length]    ,					self.vocab_size      )

          SCREAMING_SNAKE_CASE_				:      str					=	None
          if self.use_input_mask:
                SCREAMING_SNAKE_CASE_				:      Dict					=	random_attention_mask([self.batch_size, self.seq_length]      )

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

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

          SCREAMING_SNAKE_CASE_				:      List[Any]					=	None
          SCREAMING_SNAKE_CASE_				:      str					=	None
          SCREAMING_SNAKE_CASE_				:      int					=	None
          if self.use_labels:
                SCREAMING_SNAKE_CASE_				:      List[str]					=	ids_tensor([self.batch_size]    ,					self.type_sequence_label_size      )
                SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	ids_tensor([self.batch_size, self.seq_length]    ,					self.num_labels      )
                SCREAMING_SNAKE_CASE_				:      Any					=	ids_tensor([self.batch_size]    ,					self.num_choices      )

          SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	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__    ,					)

          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	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   __lowerCamelCase				( self    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					*lowercase__      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      Tuple					=	TFGPTJModel(config=lowerCamelCase__      )
          SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	{
              "input_ids": input_ids,
              "attention_mask": input_mask,
              "token_type_ids": token_type_ids,
          }
          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	model(lowerCamelCase__      )

          SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	[input_ids, None, input_mask]  # None is the input for 'past'
          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	model(lowerCamelCase__      )

          SCREAMING_SNAKE_CASE_				:      List[str]					=	model(lowerCamelCase__      )

          self.parent.assertEqual(result.last_hidden_state.shape    ,					(self.batch_size, self.seq_length, self.hidden_size)      )




    def   __lowerCamelCase				( self    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					*lowercase__      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	TFGPTJModel(config=lowerCamelCase__      )

          # first forward pass
          SCREAMING_SNAKE_CASE_				:      Tuple					=	model(lowerCamelCase__    ,					token_type_ids=lowerCamelCase__    ,					use_cache=lowerCamelCase__      )
          SCREAMING_SNAKE_CASE_				:      str					=	model(lowerCamelCase__    ,					token_type_ids=lowerCamelCase__      )
          SCREAMING_SNAKE_CASE_				:      Dict					=	model(lowerCamelCase__    ,					token_type_ids=lowerCamelCase__    ,					use_cache=lowerCamelCase__      )

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

          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	outputs.to_tuple()

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

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

          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	model(lowerCamelCase__    ,					token_type_ids=lowerCamelCase__      )["last_hidden_state"]
          SCREAMING_SNAKE_CASE_				:      Any					=	model(lowerCamelCase__    ,					token_type_ids=lowerCamelCase__    ,					past_key_values=lowerCamelCase__      )[
              "last_hidden_state"
          ]

          # select random slice
          SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	int(ids_tensor((1,)    ,					shape_list(lowerCamelCase__      )[-1]      )      )
          SCREAMING_SNAKE_CASE_				:      int					=	output_from_no_past[:, -1, random_slice_idx]
          SCREAMING_SNAKE_CASE_				:      Tuple					=	output_from_past[:, 0, random_slice_idx]

          # test that outputs are equal for slice
          tf.debugging.assert_near(lowerCamelCase__    ,					lowerCamelCase__    ,					rtol=1e-6      )




    def   __lowerCamelCase				( self    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					*lowercase__      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      str					=	TFGPTJModel(config=lowerCamelCase__      )

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

          # first forward pass
          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	model(lowerCamelCase__    ,					attention_mask=lowerCamelCase__      ).to_tuple()

          # create hypothetical next token and extent to next_input_ids
          SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	ids_tensor((self.batch_size, 1)    ,					config.vocab_size      )

          # change a random masked slice from input_ids
          SCREAMING_SNAKE_CASE_				:      List[Any]					=	ids_tensor((1,)    ,					lowerCamelCase__      ).numpy() + 1
          SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	ids_tensor((self.batch_size, self.seq_length)    ,					config.vocab_size      )
          SCREAMING_SNAKE_CASE_				:      Dict					=	tf.range(self.seq_length      ) == (self.seq_length - random_seq_idx_to_change)
          SCREAMING_SNAKE_CASE_				:      Tuple					=	tf.transpose(
              tf.broadcast_to(tf.expand_dims(lowerCamelCase__    ,					-1      )    ,					(self.seq_length, self.batch_size)      )      )
          SCREAMING_SNAKE_CASE_				:      Tuple					=	tf.where(lowerCamelCase__    ,					lowerCamelCase__    ,					lowerCamelCase__      )

          # append to next input_ids and attn_mask
          SCREAMING_SNAKE_CASE_				:      str					=	tf.concat([input_ids, next_tokens]    ,					axis=-1      )
          SCREAMING_SNAKE_CASE_				:      Tuple					=	tf.concat([attn_mask, tf.ones((shape_list(lowerCamelCase__      )[0], 1)    ,					dtype=tf.intaa      )]    ,					axis=1      )

          # get two different outputs
          SCREAMING_SNAKE_CASE_				:      List[Any]					=	model(lowerCamelCase__    ,					attention_mask=lowerCamelCase__      )["last_hidden_state"]
          SCREAMING_SNAKE_CASE_				:      Dict					=	model(lowerCamelCase__    ,					past_key_values=lowerCamelCase__    ,					attention_mask=lowerCamelCase__      )[
              "last_hidden_state"
          ]

          # select random slice
          SCREAMING_SNAKE_CASE_				:      Tuple					=	int(ids_tensor((1,)    ,					shape_list(lowerCamelCase__      )[-1]      )      )
          SCREAMING_SNAKE_CASE_				:      List[Any]					=	output_from_no_past[:, -1, random_slice_idx]
          SCREAMING_SNAKE_CASE_				:      Any					=	output_from_past[:, 0, random_slice_idx]

          # test that outputs are equal for slice
          tf.debugging.assert_near(lowerCamelCase__    ,					lowerCamelCase__    ,					rtol=1e-12      )




    def   __lowerCamelCase				( self    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					*lowercase__      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      Any					=	TFGPTJModel(config=lowerCamelCase__      )

          SCREAMING_SNAKE_CASE_				:      int					=	input_ids[:1, :]
          SCREAMING_SNAKE_CASE_				:      Dict					=	input_mask[:1, :]
          SCREAMING_SNAKE_CASE_				:      Dict					=	token_type_ids[:1, :]
          SCREAMING_SNAKE_CASE_				:      List[str]					=	1

          # first forward pass
          SCREAMING_SNAKE_CASE_				:      Dict					=	model(lowerCamelCase__    ,					attention_mask=lowerCamelCase__    ,					token_type_ids=lowerCamelCase__    ,					use_cache=lowerCamelCase__      )

          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	outputs.to_tuple()

          # create hypothetical next token and extent to next_input_ids
          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	ids_tensor((self.batch_size, 3)    ,					config.vocab_size      )
          SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	ids_tensor((self.batch_size, 3)    ,					2      )
          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	ids_tensor((self.batch_size, 3)    ,					self.type_vocab_size      )

          # append to next input_ids and token_type_ids
          SCREAMING_SNAKE_CASE_				:      List[str]					=	tf.concat([input_ids, next_tokens]    ,					axis=-1      )
          SCREAMING_SNAKE_CASE_				:      Tuple					=	tf.concat([input_mask, next_attn_mask]    ,					axis=-1      )
          SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	tf.concat([token_type_ids, next_token_types]    ,					axis=-1      )

          SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	model(
              lowerCamelCase__    ,					token_type_ids=lowerCamelCase__    ,					attention_mask=lowerCamelCase__      )["last_hidden_state"]
          SCREAMING_SNAKE_CASE_				:      Any					=	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
          SCREAMING_SNAKE_CASE_				:      List[str]					=	int(ids_tensor((1,)    ,					shape_list(lowerCamelCase__      )[-1]      )      )
          SCREAMING_SNAKE_CASE_				:      Tuple					=	output_from_no_past[:, -3:, random_slice_idx]
          SCREAMING_SNAKE_CASE_				:      int					=	output_from_past[:, :, random_slice_idx]

          # test that outputs are equal for slice
          tf.debugging.assert_near(lowerCamelCase__    ,					lowerCamelCase__    ,					rtol=1e-3      )




    def   __lowerCamelCase				( self    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					*lowercase__      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	TFGPTJForCausalLM(config=lowerCamelCase__      )
          SCREAMING_SNAKE_CASE_				:      List[Any]					=	{
              "input_ids": input_ids,
              "attention_mask": input_mask,
              "token_type_ids": token_type_ids,
          }
          SCREAMING_SNAKE_CASE_				:      Any					=	model(lowerCamelCase__      )
          self.parent.assertEqual(result.logits.shape    ,					(self.batch_size, self.seq_length, self.vocab_size)      )







    def   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	self.prepare_config_and_inputs()

          (
              (
              SCREAMING_SNAKE_CASE_  
          ),						(
              SCREAMING_SNAKE_CASE_  
          ),						(
              SCREAMING_SNAKE_CASE_  
          ),						(
              SCREAMING_SNAKE_CASE_  
          ),						(
              SCREAMING_SNAKE_CASE_  
          ),						(
              SCREAMING_SNAKE_CASE_  
          ),						(
              SCREAMING_SNAKE_CASE_  
          ),						(
              SCREAMING_SNAKE_CASE_  
          ),						(
              SCREAMING_SNAKE_CASE_  
          ),						  
          )				:      str					=	config_and_inputs

          SCREAMING_SNAKE_CASE_				:      List[Any]					=	{
              "input_ids": input_ids,
              "token_type_ids": token_type_ids,
              "attention_mask": input_mask,
          }
          return config, inputs_dict






@require_tf
class      SCREAMING_SNAKE_CASE__      (    __UpperCAmelCase,__UpperCAmelCase,__UpperCAmelCase,unittest.TestCase				):
    _A         =							(
        (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
        if is_tf_available()
        else ()
    )

    _A         =							(TFGPTJForCausalLM,) if is_tf_available() else ()
    _A         =							(
        {
            """feature-extraction""": TFGPTJModel,
            """question-answering""": TFGPTJForQuestionAnswering,
            """text-classification""": TFGPTJForSequenceClassification,
            """text-generation""": TFGPTJForCausalLM,
            """zero-shot""": TFGPTJForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    _A         =							False
    _A         =							False
    _A         =							False
    _A         =							False




    def   __lowerCamelCase				( self    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__      ):






          """simple docstring"""






          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   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      str					=	TFGPTJModelTester(self      )
          SCREAMING_SNAKE_CASE_				:      Any					=	ConfigTester(self    ,					config_class=lowerCamelCase__    ,					n_embd=37      )




    def   __lowerCamelCase				( self      ):






          """simple docstring"""






          self.config_tester.run_common_tests()




    def   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_gptj_model(*lowerCamelCase__      )




    def   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_gptj_model_past(*lowerCamelCase__      )




    def   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_gptj_model_attention_mask_past(*lowerCamelCase__      )




    def   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      List[Any]					=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_gptj_model_past_large_inputs(*lowerCamelCase__      )




    def   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      List[str]					=	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   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      int					=	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   __lowerCamelCase				( self      ):






          """simple docstring"""






          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      SCREAMING_SNAKE_CASE__      (    unittest.TestCase				):




    def   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      str					=	TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"    ,					from_pt=lowerCamelCase__      )
          SCREAMING_SNAKE_CASE_				:      Tuple					=	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
          SCREAMING_SNAKE_CASE_				:      int					=	[464, 3290, 318, 257, 582, 338, 1266, 1545, 13, 632, 318, 257, 9112, 1_5185, 11, 290, 340, 318, 257, 1545]
          # fmt: on
          SCREAMING_SNAKE_CASE_				:      List[str]					=	model.generate(lowerCamelCase__    ,					do_sample=lowerCamelCase__      )
          self.assertListEqual(output_ids[0].numpy().tolist()    ,					lowerCamelCase__      )




    def   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      str					=	AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"    ,					revision="float16"      )
          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"    ,					revision="float16"    ,					from_pt=lowerCamelCase__      )

          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	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"      ):
                SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	model.generate(**lowerCamelCase__    ,					do_sample=lowerCamelCase__    ,					seed=[42, 0]      )
          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	tokenizer.decode(output_ids[0]    ,					skip_special_tokens=lowerCamelCase__      )

          SCREAMING_SNAKE_CASE_				:      List[str]					=	"Today is a nice day and I’m going to go for a walk. I’"
          self.assertEqual(lowerCamelCase__    ,					lowerCamelCase__      )




    def   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      List[str]					=	TFGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"    ,					revision="float16"    ,					from_pt=lowerCamelCase__      )
          SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"    ,					revision="float16"      )

          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	"left"

          # Define PAD Token = EOS Token = 50256
          SCREAMING_SNAKE_CASE_				:      Dict					=	tokenizer.eos_token
          SCREAMING_SNAKE_CASE_				:      List[str]					=	model.config.eos_token_id

          # use different length sentences to test batching
          SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	[
              "Hello, my dog is a little",
              "Today, I",
          ]
          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	[
              "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   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_				:      str					=	self._get_beam_search_test_objects()

          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	tokenizer(lowerCamelCase__    ,					return_tensors="tf"    ,					padding=lowerCamelCase__      )
          SCREAMING_SNAKE_CASE_				:      Any					=	model.generate(**lowerCamelCase__    ,					do_sample=lowerCamelCase__    ,					num_beams=2      )
          SCREAMING_SNAKE_CASE_				:      str					=	tokenizer.batch_decode(lowerCamelCase__    ,					skip_special_tokens=lowerCamelCase__      )
          self.assertListEqual(lowerCamelCase__    ,					lowerCamelCase__      )




    def   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_				:      Tuple					=	self._get_beam_search_test_objects()

          SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	tokenizer(lowerCamelCase__    ,					return_tensors="tf"    ,					padding=lowerCamelCase__      )
          SCREAMING_SNAKE_CASE_				:      str					=	tokenizer(sentences[0]    ,					return_tensors="tf"      )
          SCREAMING_SNAKE_CASE_				:      Union[str, Any]					=	model.generate(**lowerCamelCase__    ,					do_sample=lowerCamelCase__    ,					num_beams=2      )
          SCREAMING_SNAKE_CASE_				:      int					=	(
              shape_list(inputs_non_padded["input_ids"]      )[-1]
              - tf.reduce_sum(tf.cast(inputs["attention_mask"][-1]    ,					tf.intaa      )      ).numpy()
          )
          SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	tokenizer(sentences[1]    ,					return_tensors="tf"      )
          SCREAMING_SNAKE_CASE_				:      List[str]					=	model.generate(
              **lowerCamelCase__    ,					do_sample=lowerCamelCase__    ,					num_beams=2    ,					max_length=model.config.max_length - num_paddings      )
          SCREAMING_SNAKE_CASE_				:      List[Any]					=	tokenizer.decode(output_non_padded[0]    ,					skip_special_tokens=lowerCamelCase__      )
          SCREAMING_SNAKE_CASE_				:      Tuple					=	tokenizer.decode(output_padded[0]    ,					skip_special_tokens=lowerCamelCase__      )
          self.assertListEqual(lowerCamelCase__    ,					[non_padded_sentence, padded_sentence]      )







    def   __lowerCamelCase				( self      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_				:      List[Any]					=	self._get_beam_search_test_objects()

          SCREAMING_SNAKE_CASE_				:      Tuple					=	tokenizer(lowerCamelCase__    ,					return_tensors="tf"    ,					padding=lowerCamelCase__      )
          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	tf.function(model.generate    ,					jit_compile=lowerCamelCase__      )
          SCREAMING_SNAKE_CASE_				:      Tuple					=	xla_generate(**lowerCamelCase__    ,					do_sample=lowerCamelCase__    ,					num_beams=2      )
          SCREAMING_SNAKE_CASE_				:      List[str]					=	tokenizer.batch_decode(lowerCamelCase__    ,					skip_special_tokens=lowerCamelCase__      )
          self.assertListEqual(lowerCamelCase__    ,					lowerCamelCase__      )




'''simple docstring'''






import os
from argparse import ArgumentParser
from pathlib import Path
from shutil import copyfile
from typing import List

from datasets import config
from datasets.builder import DatasetBuilder
from datasets.commands import BaseDatasetsCLICommand
from datasets.download.download_config import DownloadConfig
from datasets.download.download_manager import DownloadMode
from datasets.load import dataset_module_factory, import_main_class
from datasets.utils.info_utils import VerificationMode


def  __lowerCamelCase      (				SCREAMING_SNAKE_CASE_					:     int  ,  **SCREAMING_SNAKE_CASE_					:     int   )							->       str:
      """simple docstring"""

      return RunBeamCommand(
          args.dataset  ,  args.name  ,  args.cache_dir  ,  args.beam_pipeline_options  ,  args.data_dir  ,  args.all_configs  ,  args.save_info or args.save_infos  ,  args.ignore_verifications  ,  args.force_redownload  ,  **SCREAMING_SNAKE_CASE_  ,  )





class      SCREAMING_SNAKE_CASE__      (    _UpperCAmelCase				):




    @staticmethod
    def   __lowerCamelCase				( lowercase__      ):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      str					=	parser.add_parser("run_beam"    ,					help="Run a Beam dataset processing pipeline"      )
          run_beam_parser.add_argument("dataset"    ,					type=lowercase__    ,					help="Name of the dataset to download"      )
          run_beam_parser.add_argument("--name"    ,					type=lowercase__    ,					default=lowercase__    ,					help="Dataset config name"      )
          run_beam_parser.add_argument(
              "--cache_dir"    ,					type=lowercase__    ,					default=lowercase__    ,					help="Cache directory where the datasets are stored"    ,					)
          run_beam_parser.add_argument(
              "--beam_pipeline_options"    ,					type=lowercase__    ,					default=""    ,					help="Beam pipeline options, separated by commas. Example:: `--beam_pipeline_options=job_name=my-job,project=my-project`"    ,					)
          run_beam_parser.add_argument(
              "--data_dir"    ,					type=lowercase__    ,					default=lowercase__    ,					help="Can be used to specify a manual directory to get the files from"    ,					)
          run_beam_parser.add_argument("--all_configs"    ,					action="store_true"    ,					help="Test all dataset configurations"      )
          run_beam_parser.add_argument("--save_info"    ,					action="store_true"    ,					help="Save the dataset infos file"      )
          run_beam_parser.add_argument(
              "--ignore_verifications"    ,					action="store_true"    ,					help="Run the test without checksums and splits checks"      )
          run_beam_parser.add_argument("--force_redownload"    ,					action="store_true"    ,					help="Force dataset redownload"      )
          # aliases
          run_beam_parser.add_argument("--save_infos"    ,					action="store_true"    ,					help="alias for save_info"      )
          run_beam_parser.set_defaults(func=lowercase__      )




    def __init__( self    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__    ,					**lowercase__    ,					):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_				:      Tuple					=	dataset
          SCREAMING_SNAKE_CASE_				:      Dict					=	name
          SCREAMING_SNAKE_CASE_				:      str					=	cache_dir
          SCREAMING_SNAKE_CASE_				:      Optional[Any]					=	beam_pipeline_options
          SCREAMING_SNAKE_CASE_				:      List[Any]					=	data_dir
          SCREAMING_SNAKE_CASE_				:      Dict					=	all_configs
          SCREAMING_SNAKE_CASE_				:      List[str]					=	save_infos
          SCREAMING_SNAKE_CASE_				:      List[Any]					=	ignore_verifications
          SCREAMING_SNAKE_CASE_				:      int					=	force_redownload
          SCREAMING_SNAKE_CASE_				:      Tuple					=	config_kwargs







    def   __lowerCamelCase				( self      ):






          """simple docstring"""






          import apache_beam as beam

          if self._name is not None and self._all_configs:
                print("Both parameters `name` and `all_configs` can't be used at once."      )
                exit(1      )
          SCREAMING_SNAKE_CASE_,						SCREAMING_SNAKE_CASE_				:      Any					=	self._dataset, self._name
          SCREAMING_SNAKE_CASE_				:      List[str]					=	dataset_module_factory(lowercase__      )
          SCREAMING_SNAKE_CASE_				:      Optional[int]					=	import_main_class(dataset_module.module_path      )
          SCREAMING_SNAKE_CASE_				:      List[DatasetBuilder]					=	[]
          if self._beam_pipeline_options:
                SCREAMING_SNAKE_CASE_				:      List[str]					=	beam.options.pipeline_options.PipelineOptions(
                    flags=[F"--{opt.strip()}" for opt in self._beam_pipeline_options.split(","      ) if opt]      )
          else:
                SCREAMING_SNAKE_CASE_				:      Dict					=	None
          if self._all_configs and len(builder_cls.BUILDER_CONFIGS      ) > 0:
                for builder_config in builder_cls.BUILDER_CONFIGS:
                      builders.append(
                          builder_cls(
                              config_name=builder_config.name    ,					data_dir=self._data_dir    ,					hash=dataset_module.hash    ,					beam_options=lowercase__    ,					cache_dir=self._cache_dir    ,					base_path=dataset_module.builder_kwargs.get("base_path"      )    ,					)      )
          else:
                builders.append(
                    builder_cls(
                        config_name=lowercase__    ,					data_dir=self._data_dir    ,					beam_options=lowercase__    ,					cache_dir=self._cache_dir    ,					base_path=dataset_module.builder_kwargs.get("base_path"      )    ,					**self._config_kwargs    ,					)      )

          for builder in builders:
                builder.download_and_prepare(
                    download_mode=DownloadMode.REUSE_CACHE_IF_EXISTS
                    if not self._force_redownload
                    else DownloadMode.FORCE_REDOWNLOAD    ,					download_config=DownloadConfig(cache_dir=config.DOWNLOADED_DATASETS_PATH      )    ,					verification_mode=VerificationMode.NO_CHECKS
                    if self._ignore_verifications
                    else VerificationMode.ALL_CHECKS    ,					try_from_hf_gcs=lowercase__    ,					)
                if self._save_infos:
                      builder._save_infos()

          print("Apache beam run successful."      )

          # If save_infos=True, the dataset infos file is created next to the loaded module file.
          # Let's move it to the original directory of the dataset script, to allow the user to
          # upload them on S3 at the same time afterwards.
          if self._save_infos:
                SCREAMING_SNAKE_CASE_				:      Dict					=	os.path.join(builder_cls.get_imported_module_dir()    ,					config.DATASETDICT_INFOS_FILENAME      )

                SCREAMING_SNAKE_CASE_				:      List[Any]					=	Path(lowercase__      ).name + ".py"

                SCREAMING_SNAKE_CASE_				:      Optional[int]					=	os.path.join(lowercase__    ,					lowercase__      )
                if os.path.isfile(lowercase__      ):
                      SCREAMING_SNAKE_CASE_				:      Any					=	os.path.dirname(lowercase__      )
                elif os.path.isfile(lowercase__      ):
                      SCREAMING_SNAKE_CASE_				:      List[str]					=	path
                else:  # in case of a remote dataset
                      print(F"Dataset Infos file saved at {dataset_infos_path}"      )
                      exit(1      )

                # Move datasetinfo back to the user
                SCREAMING_SNAKE_CASE_				:      Optional[int]					=	os.path.join(lowercase__    ,					config.DATASETDICT_INFOS_FILENAME      )
                copyfile(lowercase__    ,					lowercase__      )
                print(F"Dataset Infos file saved at {user_dataset_infos_path}"      )

