



from __future__ import annotations

import os
import tempfile
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import is_tensorflow_text_available, is_tf_available
from transformers.testing_utils import require_tensorflow_text, require_tf, slow

from ..test_modeling_tf_common import floats_tensor
from .test_framework_agnostic import GenerationIntegrationTestsMixin


if is_tf_available():
							import tensorflow as tf

							from transformers import (
							    AutoTokenizer,
							    TFAutoModelForCausalLM,
							    TFAutoModelForSeqaSeqLM,
							    TFAutoModelForSpeechSeqaSeq,
							    TFAutoModelForVisionaSeq,
							    TFBartForConditionalGeneration,
							    TFLogitsProcessorList,
							    TFMinLengthLogitsProcessor,
							    tf_top_k_top_p_filtering,
							)

if is_tensorflow_text_available():
							import tensorflow_text as text







@require_tf
class 			__lowerCAmelCase    (			unittest.TestCase						):



			def     _lowerCamelCase       (						self   :  Optional[int])   ->    Dict:



						"""simple docstring"""






						_UpperCAmelCase		       =   tf.convert_to_tensor(
						    [
						        [
						            8.2_2_2_0_9_9_1,  # 3rd highest value; idx. 0
						            -0.5_6_2_0_0_4_4,
						            5.2_3_2_2_9_7_5_2,
						            4.0_3_8_6_3_9_3,
						            -6.8_7_9_8_3_7_8,
						            -0.5_4_7_8_5_8_0_2,
						            -3.2_0_1_2_1_5_3,
						            2.9_2_7_7_7_1_7_6,
						            1.8_8_1_7_1_9_5_3,
						            7.3_5_3_4_1_2_7_6,  # 5th highest value; idx. 9
						            8.4_3_2_0_7_8_3_3,  # 2nd highest value; idx. 10
						            -9.8_5_7_1_1_8_3_6,
						            -5.9_6_2_0_9_2_3_6,
						            -1.1_3_0_3_9_1_6_1,
						            -7.1_1_1_5_2_9_4,
						            -0.8_3_6_9_6_3_3,
						            -5.3_1_8_6_4_0_8,
						            7.0_6_4_2_7_4_0_7,
						            0.8_1_3_6_9_3_4_4,
						            -0.8_2_0_2_3_8_1_7,
						            -5.9_1_7_9_7_9_6,
						            0.5_8_8_1_3_4_4_3,
						            -6.9_9_7_7_8_4_3_8,
						            4.7_1_5_5_1_1_8_9,
						            -0.1_8_7_7_1_6_3_7,
						            7.4_4_0_2_0_7_5_9,  # 4th highest value; idx. 25
						            9.3_8_4_5_0_9_8_7,  # 1st highest value; idx. 26
						            2.1_2_6_6_2_9_4_1,
						            -9.3_2_5_6_2_0_3_8,
						            2.3_5_6_5_2_5_2_2,
						        ],  # cummulative prob of 5 highest values <= 0.6
						        [
						            0.5_8_4_2_5_5_1_8,
						            4.5_3_1_3_9_2_3_8,
						            -5.5_7_5_1_0_4_6_4,
						            -6.2_8_0_3_0_6_9_9,
						            -7.1_9_5_2_9_5_0_3,
						            -4.0_2_1_2_2_5_5_1,
						            1.3_9_3_3_7_0_3_7,
						            -6.0_6_7_0_7_0_5_7,
						            1.5_9_4_8_0_5_1_7,
						            -9.6_4_3_1_1_9,
						            0.0_3_9_0_7_7_9_9,
						            0.6_7_2_3_1_7_6_2,
						            -8.8_8_2_0_6_7_2_6,
						            6.2_7_1_1_5_9_2_2,  # 4th highest value; idx. 13
						            2.2_8_5_2_0_7_2_3,
						            4.8_2_7_6_7_5_0_6,
						            4.3_0_4_2_1_3_6_8,
						            8.8_2_7_5_3_1_3,  # 2nd highest value; idx. 17
						            5.4_4_0_2_9_9_5_8,  # 5th highest value; idx. 18
						            -4.4_7_3_5_7_9_4,
						            7.3_8_5_7_9_5_3_6,  # 3rd highest value; idx. 20
						            -2.9_1_0_5_1_6_6_3,
						            2.6_1_9_4_6_0_7_7,
						            -2.5_6_7_4_7_6_2,
						            -9.4_8_9_5_9_3_0_2,
						            -4.0_2_9_2_2_6_4_5,
						            -1.3_5_4_1_6_9_1_8,
						            9.6_7_7_0_2_3_2_3,  # 1st highest value; idx. 27
						            -5.8_9_4_7_8_5_5_3,
						            1.8_5_3_7_0_4_6_7,
						        ],  # cummulative prob of 5 highest values <= 0.6
						    ] ,		dtype=tf.floataa ,		)

						_UpperCAmelCase		       =   tf.convert_to_tensor(
						    [[0, 0], [0, 9], [0, 10], [0, 25], [0, 26], [1, 13], [1, 17], [1, 18], [1, 20], [1, 27]] ,		dtype=tf.intaa ,		)  # expected non filtered idx as noted above

						_UpperCAmelCase		       =   tf.convert_to_tensor(
						    [8.2_2_2_0_9_9, 7.3_5_3_4_1_2_6, 8.4_3_2_0_7_8, 7.4_4_0_2_0_7_5, 9.3_8_4_5_1, 6.2_7_1_1_5_9, 8.8_2_7_5_3_1, 5.4_4_0_2_9_9_5, 7.3_8_5_7_9_5_6, 9.6_7_7_0_2_3] ,		dtype=tf.floataa ,		)  # expected non filtered values as noted above

						_UpperCAmelCase		       =   tf_top_k_top_p_filtering(__lowerCAmelCase ,		top_k=10 ,		top_p=0.6 ,		min_tokens_to_keep=4)

						_UpperCAmelCase		       =   output[output != -float('inf')]
						_UpperCAmelCase		       =   tf.cast(
						    tf.where(tf.not_equal(__lowerCAmelCase ,		tf.constant(-float('inf') ,		dtype=tf.floataa))) ,		dtype=tf.intaa ,		)

						tf.debugging.assert_near(__lowerCAmelCase ,		__lowerCAmelCase ,		rtol=1E-12)
						tf.debugging.assert_equal(__lowerCAmelCase ,		__lowerCAmelCase)







@require_tf
class 			__lowerCAmelCase    (			unittest.TestCase				,						a__						):



			if is_tf_available():
						UpperCamelCase   					=     {
						    '''AutoModelForCausalLM''': TFAutoModelForCausalLM,
						    '''AutoModelForSpeechSeq2Seq''': TFAutoModelForSpeechSeqaSeq,
						    '''AutoModelForSeq2SeqLM''': TFAutoModelForSeqaSeqLM,
						    '''AutoModelForVision2Seq''': TFAutoModelForVisionaSeq,
						    '''LogitsProcessorList''': TFLogitsProcessorList,
						    '''MinLengthLogitsProcessor''': TFMinLengthLogitsProcessor,
						    '''create_tensor_fn''': tf.convert_to_tensor,
						    '''floats_tensor''': floats_tensor,
						    '''return_tensors''': '''tf''',
						}



			@slow
			def     _lowerCamelCase       (						self   :  Tuple)   ->    Any:



						"""simple docstring"""






						_UpperCAmelCase		       =   TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2')
						_UpperCAmelCase		       =   2
						_UpperCAmelCase		       =   2







						class 			__lowerCAmelCase    (			tf.Module						):



									def __init__(						self   :  Union[str, Any] ,		A   :  Union[str, Any])   ->    Union[str, Any]:



												"""simple docstring"""






												super(__lowerCAmelCase ,		self).__init__()
												_UpperCAmelCase		       =   model



									@tf.function(
									    input_signature=(
									        tf.TensorSpec((None, input_length) ,		tf.intaa ,		name='input_ids'),
									        tf.TensorSpec((None, input_length) ,		tf.intaa ,		name='attention_mask'),
									    ) ,		jit_compile=__lowerCAmelCase ,		)
									def     _lowerCamelCase       (						self   :  int ,		A   :  List[str] ,		A   :  Optional[int])   ->    str:



												"""simple docstring"""






												_UpperCAmelCase		       =   self.model.generate(
												    input_ids=__lowerCAmelCase ,		attention_mask=__lowerCAmelCase ,		max_new_tokens=__lowerCAmelCase ,		return_dict_in_generate=__lowerCAmelCase ,		)
												return {"sequences": outputs["sequences"]}





						_UpperCAmelCase		       =   [[2, 0], [1_02, 1_03]]
						_UpperCAmelCase		       =   [[1, 0], [1, 1]]
						_UpperCAmelCase		       =   DummyModel(model=__lowerCAmelCase)
						with tempfile.TemporaryDirectory() as tmp_dir:
									tf.saved_model.save(__lowerCAmelCase ,		__lowerCAmelCase ,		signatures={'serving_default': dummy_model.serving})
									_UpperCAmelCase		       =   tf.saved_model.load(__lowerCAmelCase).signatures['''serving_default''']
									for batch_size in range(1 ,		len(__lowerCAmelCase) + 1):
												_UpperCAmelCase		       =   {
												    '''input_ids''': tf.constant(dummy_input_ids[:batch_size]),
												    '''attention_mask''': tf.constant(dummy_attention_masks[:batch_size]),
												}
												_UpperCAmelCase		       =   serving_func(**__lowerCAmelCase)['''sequences''']
												_UpperCAmelCase		       =   test_model.generate(**__lowerCAmelCase ,		max_new_tokens=__lowerCAmelCase)
												tf.debugging.assert_equal(__lowerCAmelCase ,		__lowerCAmelCase)



			@slow
			def     _lowerCamelCase       (						self   :  str)   ->    Optional[Any]:



						"""simple docstring"""






						_UpperCAmelCase		       =   TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2')
						_UpperCAmelCase		       =   1
						_UpperCAmelCase		       =   2







						class 			__lowerCAmelCase    (			tf.Module						):



									def __init__(						self   :  Optional[int] ,		A   :  int)   ->    int:



												"""simple docstring"""






												super(__lowerCAmelCase ,		self).__init__()
												_UpperCAmelCase		       =   model



									@tf.function(
									    input_signature=(
									        tf.TensorSpec((batch_size, None) ,		tf.intaa ,		name='input_ids'),
									        tf.TensorSpec((batch_size, None) ,		tf.intaa ,		name='attention_mask'),
									    ) ,		jit_compile=__lowerCAmelCase ,		)
									def     _lowerCamelCase       (						self   :  str ,		A   :  str ,		A   :  Union[str, Any])   ->    Tuple:



												"""simple docstring"""






												_UpperCAmelCase		       =   self.model.generate(
												    input_ids=__lowerCAmelCase ,		attention_mask=__lowerCAmelCase ,		max_new_tokens=__lowerCAmelCase ,		return_dict_in_generate=__lowerCAmelCase ,		)
												return {"sequences": outputs["sequences"]}





						_UpperCAmelCase		       =   [[2], [1_02, 1_03]]
						_UpperCAmelCase		       =   [[1], [1, 1]]
						_UpperCAmelCase		       =   DummyModel(model=__lowerCAmelCase)
						with tempfile.TemporaryDirectory() as tmp_dir:
									tf.saved_model.save(__lowerCAmelCase ,		__lowerCAmelCase ,		signatures={'serving_default': dummy_model.serving})
									_UpperCAmelCase		       =   tf.saved_model.load(__lowerCAmelCase).signatures['''serving_default''']
									for input_row in range(len(__lowerCAmelCase)):
												_UpperCAmelCase		       =   {
												    '''input_ids''': tf.constant([dummy_input_ids[input_row]]),
												    '''attention_mask''': tf.constant([dummy_attention_masks[input_row]]),
												}
												_UpperCAmelCase		       =   serving_func(**__lowerCAmelCase)['''sequences''']
												_UpperCAmelCase		       =   test_model.generate(**__lowerCAmelCase ,		max_new_tokens=__lowerCAmelCase)
												tf.debugging.assert_equal(__lowerCAmelCase ,		__lowerCAmelCase)



			@slow
			@require_tensorflow_text
			def     _lowerCamelCase       (						self   :  Optional[Any])   ->    List[Any]:



						"""simple docstring"""






						with tempfile.TemporaryDirectory() as tmp_dir:
									# file needed to load the TF tokenizer
									hf_hub_download(repo_id='google/flan-t5-small' ,		filename='spiece.model' ,		local_dir=__lowerCAmelCase)







									class 			__lowerCAmelCase    (			tf.keras.layers.Layer						):



												def __init__(						self   :  List[str])   ->    Any:



															"""simple docstring"""






															super().__init__()
															_UpperCAmelCase		       =   text.SentencepieceTokenizer(
															    model=tf.io.gfile.GFile(os.path.join(__lowerCAmelCase ,		'spiece.model') ,		'rb').read())
															_UpperCAmelCase		       =   TFAutoModelForSeqaSeqLM.from_pretrained('hf-internal-testing/tiny-random-t5')



												def     _lowerCamelCase       (						self   :  Union[str, Any] ,		A   :  Optional[Any] ,		*A   :  Any ,		**A   :  Tuple)   ->    List[str]:



															"""simple docstring"""






															_UpperCAmelCase		       =   self.tokenizer.tokenize(__lowerCAmelCase)
															_UpperCAmelCase		       =   text.pad_model_inputs(
															    __lowerCAmelCase ,		max_seq_length=64 ,		pad_value=self.model.config.pad_token_id)
															_UpperCAmelCase		       =   self.model.generate(input_ids=__lowerCAmelCase ,		attention_mask=__lowerCAmelCase)
															return self.tokenizer.detokenize(__lowerCAmelCase)





									_UpperCAmelCase		       =   CompleteSentenceTransformer()
									_UpperCAmelCase		       =   tf.keras.layers.Input(shape=(1,) ,		dtype=tf.string ,		name='inputs')
									_UpperCAmelCase		       =   complete_model(__lowerCAmelCase)
									_UpperCAmelCase		       =   tf.keras.Model(__lowerCAmelCase ,		__lowerCAmelCase)
									keras_model.save(__lowerCAmelCase)



			def     _lowerCamelCase       (						self   :  Any)   ->    Optional[int]:



						"""simple docstring"""






						_UpperCAmelCase		       =   {
						    '''do_sample''': True,
						    '''num_beams''': 1,
						    '''top_p''': 0.7,
						    '''top_k''': 10,
						    '''temperature''': 0.7,
						}
						_UpperCAmelCase		       =   14

						_UpperCAmelCase		       =   AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2')
						_UpperCAmelCase		       =   '''Hello, my dog is cute and'''
						_UpperCAmelCase		       =   tokenizer(__lowerCAmelCase ,		return_tensors='tf')
						_UpperCAmelCase		       =   TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2')

						_UpperCAmelCase		       =   6_38
						# forces the generation to happen on CPU, to avoid GPU-related quirks
						with tf.device(':/CPU:0'):
									tf.random.set_seed(0)
									_UpperCAmelCase		       =   model.generate(**__lowerCAmelCase ,		eos_token_id=__lowerCAmelCase ,		**__lowerCAmelCase)
						self.assertTrue(expectation == len(generated_tokens[0]))

						_UpperCAmelCase		       =   [6_38, 1_98]
						with tf.device(':/CPU:0'):
									tf.random.set_seed(0)
									_UpperCAmelCase		       =   model.generate(**__lowerCAmelCase ,		eos_token_id=__lowerCAmelCase ,		**__lowerCAmelCase)
						self.assertTrue(expectation == len(generated_tokens[0]))



			def     _lowerCamelCase       (						self   :  str)   ->    List[str]:



						"""simple docstring"""






						_UpperCAmelCase		       =   AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-bart')
						_UpperCAmelCase		       =   '''Hugging Face is a technology company based in New York and Paris.'''
						_UpperCAmelCase		       =   bart_tokenizer(__lowerCAmelCase ,		return_tensors='tf').input_ids
						_UpperCAmelCase		       =   TFBartForConditionalGeneration.from_pretrained('hf-internal-testing/tiny-random-bart')
						_UpperCAmelCase		       =   bart_model.generate(__lowerCAmelCase).numpy()







						class 			__lowerCAmelCase    (			a__						):



									def     _lowerCamelCase       (						self   :  Optional[Any] ,		A   :  int ,		A   :  List[Any]=None ,		**A   :  List[Any])   ->    Union[str, Any]:



												"""simple docstring"""






												return super().call(__lowerCAmelCase ,		**__lowerCAmelCase)





						_UpperCAmelCase		       =   FakeBart.from_pretrained('hf-internal-testing/tiny-random-bart')
						_UpperCAmelCase		       =   bart_model.generate(__lowerCAmelCase ,		foo='bar').numpy()
						self.assertTrue(np.array_equal(__lowerCAmelCase ,		__lowerCAmelCase))







						class 			__lowerCAmelCase    (			bart_model.model.encoder.__class__						):



									def     _lowerCamelCase       (						self   :  Union[str, Any] ,		A   :  Union[str, Any] ,		**A   :  str)   ->    Optional[int]:



												"""simple docstring"""






												return super().call(__lowerCAmelCase ,		**__lowerCAmelCase)





						_UpperCAmelCase		       =   FakeEncoder(bart_model.config ,		bart_model.model.shared)
						_UpperCAmelCase		       =   fake_encoder

						# Normal generation still works (the output will be different because the encoder weights are different)
						_UpperCAmelCase		       =   bart_model.generate(__lowerCAmelCase).numpy()
						with self.assertRaises(__lowerCAmelCase):
									# FakeEncoder.call() accepts **kwargs -> no filtering -> value error due to unexpected input "foo"
									bart_model.generate(__lowerCAmelCase ,		foo='bar')

'''simple docstring'''

import os
from argparse import ArgumentParser
from typing import List

import torch.utils.data

from datasets import Dataset, IterableDataset
from datasets.distributed import split_dataset_by_node


__a: Optional[int]          =   4
__a: Optional[Any]          =   3



class 			UpperCAmelCase  (						a__					):





       '''simple docstring'''



       pass


def 			__UpperCamelCase       (   UpperCAmelCase     ):
 for shard in shards:
  for i in range(UpperCAmelCase     ):
   yield {"i": i, "shard": shard}
def 			__UpperCamelCase       (   ):
 lowercase__				:    Tuple  =					int(os.environ['''RANK''']     )
 lowercase__				:    List[str]  =					int(os.environ['''WORLD_SIZE''']     )

 lowercase__				:    Optional[Any]  =					ArgumentParser()
 parser.add_argument('''--streaming'''		,     type=UpperCAmelCase     )
 parser.add_argument('''--local_rank'''		,     type=UpperCAmelCase     )
 parser.add_argument('''--num_workers'''		,     type=UpperCAmelCase		,     default=0     )
 lowercase__				:    List[Any]  =					parser.parse_args()
 lowercase__				:    List[str]  =					args.streaming
 lowercase__				:    str  =					args.num_workers

 lowercase__				:    Optional[int]  =					{'''shards''': [F"""shard_{shard_idx}""" for shard_idx in range(UpperCAmelCase     )]}
 lowercase__				:    Tuple  =					IterableDataset.from_generator(UpperCAmelCase		,     gen_kwargs=UpperCAmelCase     )
 if not streaming:
  lowercase__				:    int  =					Dataset.from_list(list(UpperCAmelCase     )     )

 lowercase__				:    str  =					split_dataset_by_node(UpperCAmelCase		,     rank=UpperCAmelCase		,     world_size=UpperCAmelCase     )
 lowercase__				:    Optional[int]  =					torch.utils.data.DataLoader(UpperCAmelCase		,     num_workers=UpperCAmelCase     )

 lowercase__				:    Dict  =					NUM_SHARDS * NUM_ITEMS_PER_SHARD
 lowercase__				:    Optional[Any]  =					full_size // world_size
 expected_local_size += int(rank < (full_size % world_size)     )

 lowercase__				:    Union[str, Any]  =					sum(1 for _ in dataloader     )
 if local_size != expected_local_size:
  raise FailedTestError(F"""local_size {local_size} != expected_local_size {expected_local_size}"""     )


if __name__ == "__main__":
  main()

