




'''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 __magic_name__:






						def __init__(   self						:						Dict					,   __UpperCamelCase						:						Tuple       ):



									'''simple docstring'''






									snake_case__          =						parent
									snake_case__          =						1_3
									snake_case__          =						7
									snake_case__          =						True
									snake_case__          =						True
									snake_case__          =						True
									snake_case__          =						True
									snake_case__          =						True
									snake_case__          =						9_9
									snake_case__          =						3_2
									snake_case__          =						4
									snake_case__          =						2
									snake_case__          =						4
									snake_case__          =						3_7
									snake_case__          =						"""gelu"""
									snake_case__          =						0.1
									snake_case__          =						0.1
									snake_case__          =						5_1_2
									snake_case__          =						1_6
									snake_case__          =						2
									snake_case__          =						0.02
									snake_case__          =						3
									snake_case__          =						4
									snake_case__          =						None
									snake_case__          =						self.vocab_size - 1
									snake_case__          =						self.vocab_size - 1
									snake_case__          =						self.vocab_size - 1






						def       __lowerCAmelCase(   self						:						str       ):



									'''simple docstring'''






									snake_case__          =						ids_tensor([self.batch_size, self.seq_length]					,   self.vocab_size       )

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

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

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

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

									snake_case__          =						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=__UpperCamelCase					,   )

									snake_case__          =						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						:						Union[str, Any]					,   __UpperCamelCase						:						Union[str, Any]					,   __UpperCamelCase						:						int					,   __UpperCamelCase						:						List[str]					,   __UpperCamelCase						:						List[Any]					,   __UpperCamelCase						:						Tuple					,   *__UpperCamelCase						:						List[Any]       ):



									'''simple docstring'''






									snake_case__          =						TFGPTJModel(config=__UpperCamelCase       )
									snake_case__          =						{
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """token_type_ids""": token_type_ids,
									}
									snake_case__          =						model(__UpperCamelCase       )

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

									snake_case__          =						model(__UpperCamelCase       )

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






						def       __lowerCAmelCase(   self						:						Optional[Any]					,   __UpperCamelCase						:						List[Any]					,   __UpperCamelCase						:						Dict					,   __UpperCamelCase						:						Tuple					,   __UpperCamelCase						:						Tuple					,   __UpperCamelCase						:						Union[str, Any]					,   *__UpperCamelCase						:						Optional[Any]       ):



									'''simple docstring'''






									snake_case__          =						TFGPTJModel(config=__UpperCamelCase       )

									# first forward pass
									snake_case__          =						model(__UpperCamelCase					,   token_type_ids=__UpperCamelCase					,   use_cache=__UpperCamelCase       )
									snake_case__          =						model(__UpperCamelCase					,   token_type_ids=__UpperCamelCase       )
									snake_case__          =						model(__UpperCamelCase					,   token_type_ids=__UpperCamelCase					,   use_cache=__UpperCamelCase       )

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

									snake_case__    ,			snake_case__          =						outputs.to_tuple()

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

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

									snake_case__          =						model(__UpperCamelCase					,   token_type_ids=__UpperCamelCase       )["""last_hidden_state"""]
									snake_case__          =						model(__UpperCamelCase					,   token_type_ids=__UpperCamelCase					,   past_key_values=__UpperCamelCase       )[
									    """last_hidden_state"""
									]

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

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






						def       __lowerCAmelCase(   self						:						Optional[int]					,   __UpperCamelCase						:						List[str]					,   __UpperCamelCase						:						List[str]					,   __UpperCamelCase						:						List[Any]					,   __UpperCamelCase						:						List[str]					,   __UpperCamelCase						:						Union[str, Any]					,   *__UpperCamelCase						:						Union[str, Any]       ):



									'''simple docstring'''






									snake_case__          =						TFGPTJModel(config=__UpperCamelCase       )

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

									# first forward pass
									snake_case__    ,			snake_case__          =						model(__UpperCamelCase					,   attention_mask=__UpperCamelCase       ).to_tuple()

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

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

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

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

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

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






						def       __lowerCAmelCase(   self						:						str					,   __UpperCamelCase						:						List[Any]					,   __UpperCamelCase						:						List[Any]					,   __UpperCamelCase						:						Optional[int]					,   __UpperCamelCase						:						Dict					,   __UpperCamelCase						:						Dict					,   *__UpperCamelCase						:						int       ):



									'''simple docstring'''






									snake_case__          =						TFGPTJModel(config=__UpperCamelCase       )

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

									# first forward pass
									snake_case__          =						model(__UpperCamelCase					,   attention_mask=__UpperCamelCase					,   token_type_ids=__UpperCamelCase					,   use_cache=__UpperCamelCase       )

									snake_case__    ,			snake_case__          =						outputs.to_tuple()

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

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

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

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

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






						def       __lowerCAmelCase(   self						:						Any					,   __UpperCamelCase						:						Optional[Any]					,   __UpperCamelCase						:						Dict					,   __UpperCamelCase						:						Dict					,   __UpperCamelCase						:						Any					,   __UpperCamelCase						:						List[Any]					,   *__UpperCamelCase						:						Tuple       ):



									'''simple docstring'''






									snake_case__          =						TFGPTJForCausalLM(config=__UpperCamelCase       )
									snake_case__          =						{
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """token_type_ids""": token_type_ids,
									}
									snake_case__          =						model(__UpperCamelCase       )
									self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.seq_length, self.vocab_size)       )







						def       __lowerCAmelCase(   self						:						str       ):



									'''simple docstring'''






									snake_case__          =						self.prepare_config_and_inputs()

									(
									    (
									    snake_case__
									)    ,			(
									    snake_case__
									)    ,			(
									    snake_case__
									)    ,			(
									    snake_case__
									)    ,			(
									    snake_case__
									)    ,			(
									    snake_case__
									)    ,			(
									    snake_case__
									)    ,			(
									    snake_case__
									)    ,			(
									    snake_case__
									)    ,			
									)          =						config_and_inputs

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







@require_tf
class __magic_name__(						__lowerCAmelCase    ,		__lowerCAmelCase    ,		__lowerCAmelCase    ,		unittest.TestCase      ):
						UpperCAmelCase_		:					Optional[Any]    =  (
						    (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
						    if is_tf_available()
						    else ()
						)

						UpperCAmelCase_		:					Optional[Any]    =  (TFGPTJForCausalLM,) if is_tf_available() else ()
						UpperCAmelCase_		:					Dict    =  (
						    {
						        """feature-extraction""": TFGPTJModel,
						        """question-answering""": TFGPTJForQuestionAnswering,
						        """text-classification""": TFGPTJForSequenceClassification,
						        """text-generation""": TFGPTJForCausalLM,
						        """zero-shot""": TFGPTJForSequenceClassification,
						    }
						    if is_tf_available()
						    else {}
						)
						UpperCAmelCase_		:					str    =  False
						UpperCAmelCase_		:					Optional[Any]    =  False
						UpperCAmelCase_		:					Union[str, Any]    =  False
						UpperCAmelCase_		:					Union[str, Any]    =  False






						def       __lowerCAmelCase(   self						:						Any					,   __UpperCamelCase						:						Optional[int]					,   __UpperCamelCase						:						Optional[Any]					,   __UpperCamelCase						:						Dict					,   __UpperCamelCase						:						Optional[int]					,   __UpperCamelCase						:						Union[str, Any]       ):



									'''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						:						Optional[int]       ):



									'''simple docstring'''






									snake_case__          =						TFGPTJModelTester(self       )
									snake_case__          =						ConfigTester(self					,   config_class=__UpperCamelCase					,   n_embd=3_7       )






						def       __lowerCAmelCase(   self						:						List[str]       ):



									'''simple docstring'''






									self.config_tester.run_common_tests()






						def       __lowerCAmelCase(   self						:						List[str]       ):



									'''simple docstring'''






									snake_case__          =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_gptj_model(*__UpperCamelCase       )






						def       __lowerCAmelCase(   self						:						Tuple       ):



									'''simple docstring'''






									snake_case__          =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_gptj_model_past(*__UpperCamelCase       )






						def       __lowerCAmelCase(   self						:						Optional[int]       ):



									'''simple docstring'''






									snake_case__          =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_gptj_model_attention_mask_past(*__UpperCamelCase       )






						def       __lowerCAmelCase(   self						:						Union[str, Any]       ):



									'''simple docstring'''






									snake_case__          =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_gptj_model_past_large_inputs(*__UpperCamelCase       )






						def       __lowerCAmelCase(   self						:						List[str]       ):



									'''simple docstring'''






									snake_case__          =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_gptj_lm_head_model(*__UpperCamelCase       )






						@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						:						Optional[int]       ):



									'''simple docstring'''






									snake_case__          =						TFGPTJModel.from_pretrained("""EleutherAI/gpt-j-6B"""					,   from_pt=__UpperCamelCase       )
									self.assertIsNotNone(__UpperCamelCase       )







						@unittest.skip(reason="""Currently, model embeddings are going to undergo a major refactor."""       )
						def       __lowerCAmelCase(   self						:						Any       ):



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






						def       __lowerCAmelCase(   self						:						Any       ):



									'''simple docstring'''






									snake_case__          =						TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""					,   from_pt=__UpperCamelCase       )
									snake_case__          =						tf.convert_to_tensor([[4_6_4, 3_2_9_0]]					,   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
									snake_case__          =						[4_6_4, 3_2_9_0, 3_1_8, 2_5_7, 5_8_2, 3_3_8, 1_2_6_6, 1_5_4_5, 1_3, 6_3_2, 3_1_8, 2_5_7, 9_1_1_2, 1_5_1_8_5, 1_1, 2_9_0, 3_4_0, 3_1_8, 2_5_7, 1_5_4_5]
									# fmt: on
									snake_case__          =						model.generate(__UpperCamelCase					,   do_sample=__UpperCamelCase       )
									self.assertListEqual(output_ids[0].numpy().tolist()					,   __UpperCamelCase       )






						def       __lowerCAmelCase(   self						:						Any       ):



									'''simple docstring'''






									snake_case__          =						AutoTokenizer.from_pretrained("""EleutherAI/gpt-j-6B"""					,   revision="""float16"""       )
									snake_case__          =						TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""					,   revision="""float16"""					,   from_pt=__UpperCamelCase       )

									snake_case__          =						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"""       ):
												snake_case__          =						model.generate(**__UpperCamelCase					,   do_sample=__UpperCamelCase					,   seed=[4_2, 0]       )
									snake_case__          =						tokenizer.decode(output_ids[0]					,   skip_special_tokens=__UpperCamelCase       )

									snake_case__          =						"""Today is a nice day and I’m going to go for a walk. I’"""
									self.assertEqual(__UpperCamelCase					,   __UpperCamelCase       )






						def       __lowerCAmelCase(   self						:						Optional[Any]       ):



									'''simple docstring'''






									snake_case__          =						TFGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""					,   revision="""float16"""					,   from_pt=__UpperCamelCase       )
									snake_case__          =						AutoTokenizer.from_pretrained("""EleutherAI/gpt-j-6B"""					,   revision="""float16"""       )

									snake_case__          =						"""left"""

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

									# use different length sentences to test batching
									snake_case__          =						[
									    """Hello, my dog is a little""",
									    """Today, I""",
									]
									snake_case__          =						[
									    """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						:						List[str]       ):



									'''simple docstring'''






									snake_case__    ,			snake_case__    ,			snake_case__    ,			snake_case__          =						self._get_beam_search_test_objects()

									snake_case__          =						tokenizer(__UpperCamelCase					,   return_tensors="""tf"""					,   padding=__UpperCamelCase       )
									snake_case__          =						model.generate(**__UpperCamelCase					,   do_sample=__UpperCamelCase					,   num_beams=2       )
									snake_case__          =						tokenizer.batch_decode(__UpperCamelCase					,   skip_special_tokens=__UpperCamelCase       )
									self.assertListEqual(__UpperCamelCase					,   __UpperCamelCase       )






						def       __lowerCAmelCase(   self						:						Union[str, Any]       ):



									'''simple docstring'''






									snake_case__    ,			snake_case__    ,			snake_case__    ,			snake_case__          =						self._get_beam_search_test_objects()

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







						def       __lowerCAmelCase(   self						:						Optional[int]       ):



									'''simple docstring'''






									snake_case__    ,			snake_case__    ,			snake_case__    ,			snake_case__          =						self._get_beam_search_test_objects()

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




import json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


UpperCAmelCase :				Any							=1000_0000_0000
UpperCAmelCase :				Any							=1_0000

UpperCAmelCase			,			UpperCAmelCase :				Tuple							=os.path.split(__file__)
UpperCAmelCase :				Tuple							=os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))



def  _lowerCAmelCase     (_lowerCAmelCase ,						_lowerCAmelCase):
					UpperCamelCase_			= pa.Table.from_pydict({"col": [0] * chunk_size})
					UpperCamelCase_			= pa.concat_tables([table] * (num_examples // chunk_size))
					return datasets.Dataset(_lowerCAmelCase ,						fingerprint="table_100B")


@dataclass
class    _lowercase :
				'''simple docstring'''







				lowercase__     =    42
				lowercase__     =    42
				lowercase__     =    42
				lowercase__     =    42

				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= np.random.default_rng(self.seed     )
									UpperCamelCase_			= rng.integers(low=self.low  ,			high=self.high  ,			size=self.size     ).tolist()

				def __iter__(			self     ):

									'''simple docstring'''


									return iter(self._sampled_values     )

				def __len__(			self     ):

									'''simple docstring'''


									return self.size




@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[0]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[-1]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[range(len(_lowerCAmelCase) // 2 ,						len(_lowerCAmelCase) // 2 + 10_24)]



@get_duration
def  _lowerCAmelCase     (_lowerCAmelCase):
					UpperCamelCase_			= dataset[RandIter(0 ,						len(_lowerCAmelCase) ,						10_24 ,						seed=42)]



def  _lowerCAmelCase     ():
					UpperCamelCase_			= {"num examples": SPEED_TEST_N_EXAMPLES}
					UpperCamelCase_			= (get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
					print("generating dataset")
					UpperCamelCase_			= generate_100B_dataset(num_examples=_lowerCAmelCase ,						chunk_size=_lowerCAmelCase)
					print("Functions")
					for func in functions:
										print(func.__name__)
										UpperCamelCase_			= func(_lowerCAmelCase)

					with open(_lowerCAmelCase ,						"wb") as f:
										f.write(json.dumps(_lowerCAmelCase).encode("utf-8"))


if __name__ == "__main__":  # useful to run the profiler
	benchmark_table_100B()
