



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 					_lowercase      (unittest.TestCase					):
				'''simple docstring'''







				def 					_lowerCamelCase    (       self			):

									'''simple docstring'''


									UpperCamelCase_						     =      tf.convert_to_tensor(
									    [
									        [
									            8.2_220_991,  # 3rd highest value; idx. 0
									            -0.5_620_044,
									            5.23_229_752,
									            4.0_386_393,
									            -6.8_798_378,
									            -0.54_785_802,
									            -3.2_012_153,
									            2.92_777_176,
									            1.88_171_953,
									            7.35_341_276,  # 5th highest value; idx. 9
									            8.43_207_833,  # 2nd highest value; idx. 10
									            -9.85_711_836,
									            -5.96_209_236,
									            -1.13_039_161,
									            -7.1_115_294,
									            -0.8_369_633,
									            -5.3_186_408,
									            7.06_427_407,
									            0.81_369_344,
									            -0.82_023_817,
									            -5.9_179_796,
									            0.58_813_443,
									            -6.99_778_438,
									            4.71_551_189,
									            -0.18_771_637,
									            7.44_020_759,  # 4th highest value; idx. 25
									            9.38_450_987,  # 1st highest value; idx. 26
									            2.12_662_941,
									            -9.32_562_038,
									            2.35_652_522,
									        ],  # cummulative prob of 5 highest values <= 0.6
									        [
									            0.58_425_518,
									            4.53_139_238,
									            -5.57_510_464,
									            -6.28_030_699,
									            -7.19_529_503,
									            -4.02_122_551,
									            1.39_337_037,
									            -6.06_707_057,
									            1.59_480_517,
									            -9.643_119,
									            0.03_907_799,
									            0.67_231_762,
									            -8.88_206_726,
									            6.27_115_922,  # 4th highest value; idx. 13
									            2.28_520_723,
									            4.82_767_506,
									            4.30_421_368,
									            8.8_275_313,  # 2nd highest value; idx. 17
									            5.44_029_958,  # 5th highest value; idx. 18
									            -4.4_735_794,
									            7.38_579_536,  # 3rd highest value; idx. 20
									            -2.91_051_663,
									            2.61_946_077,
									            -2.5_674_762,
									            -9.48_959_302,
									            -4.02_922_645,
									            -1.35_416_918,
									            9.67_702_323,  # 1st highest value; idx. 27
									            -5.89_478_553,
									            1.85_370_467,
									        ],  # 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.222_099, 7.3_534_126, 8.432_078, 7.4_402_075, 9.38_451, 6.271_159, 8.827_531, 5.4_402_995, 7.3_857_956, 9.677_023]   ,       dtype=tf.floataa   ,       )  # expected non filtered values as noted above

									UpperCamelCase_						     =      tf_top_k_top_p_filtering(__a   ,       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(__a   ,       tf.constant(-float("inf"			)   ,       dtype=tf.floataa			)			)			)   ,       dtype=tf.intaa   ,       )

									tf.debugging.assert_near(__a   ,       __a   ,       rtol=1e-12			)
									tf.debugging.assert_equal(__a   ,       __a			)




@require_tf
class 					_lowercase      (unittest.TestCase					,      a_					):
				'''simple docstring'''







				if is_tf_available():
								lowercase__        =  {
								    """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			):

									'''simple docstring'''


									UpperCamelCase_						     =      TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"			)
									UpperCamelCase_						     =      2
									UpperCamelCase_						     =      2


									class 					_lowercase      (tf.Module					):
													'''simple docstring'''







													def __init__(       self   ,       snake_case__			):

																		'''simple docstring'''


																		super(__a   ,       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=__a   ,       )
													def 					_lowerCamelCase    (       self   ,       snake_case__   ,       snake_case__			):

																		'''simple docstring'''


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




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

				@slow
				def 					_lowerCamelCase    (       self			):

									'''simple docstring'''


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


									class 					_lowercase      (tf.Module					):
													'''simple docstring'''







													def __init__(       self   ,       snake_case__			):

																		'''simple docstring'''


																		super(__a   ,       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=__a   ,       )
													def 					_lowerCamelCase    (       self   ,       snake_case__   ,       snake_case__			):

																		'''simple docstring'''


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




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

				@slow
				@require_tensorflow_text
				def 					_lowerCamelCase    (       self			):

									'''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=__a			)


														class 					_lowercase      (tf.keras.layers.Layer					):
																		'''simple docstring'''







																		def __init__(       self			):

																							'''simple docstring'''


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

																		def 					_lowerCamelCase    (       self   ,       snake_case__   ,       *snake_case__   ,       **snake_case__			):

																							'''simple docstring'''


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




														UpperCamelCase_						     =      CompleteSentenceTransformer()
														UpperCamelCase_						     =      tf.keras.layers.Input(shape=(1,)   ,       dtype=tf.string   ,       name="inputs"			)
														UpperCamelCase_						     =      complete_model(__a			)
														UpperCamelCase_						     =      tf.keras.Model(__a   ,       __a			)
														keras_model.save(__a			)

				def 					_lowerCamelCase    (       self			):

									'''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(__a   ,       return_tensors="tf"			)
									UpperCamelCase_						     =      TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"			)

									UpperCamelCase_						     =      638
									# 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(**__a   ,       eos_token_id=__a   ,       **__a			)
									self.assertTrue(expectation == len(generated_tokens[0]			)			)

									UpperCamelCase_						     =      [638, 198]
									with tf.device(":/CPU:0"			):
														tf.random.set_seed(0			)
														UpperCamelCase_						     =      model.generate(**__a   ,       eos_token_id=__a   ,       **__a			)
									self.assertTrue(expectation == len(generated_tokens[0]			)			)

				def 					_lowerCamelCase    (       self			):

									'''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(__a   ,       return_tensors="tf"			).input_ids
									UpperCamelCase_						     =      TFBartForConditionalGeneration.from_pretrained("hf-internal-testing/tiny-random-bart"			)
									UpperCamelCase_						     =      bart_model.generate(__a			).numpy()


									class 					_lowercase      (a_					):
													'''simple docstring'''







													def 					_lowerCamelCase    (       self   ,       snake_case__   ,       snake_case__=None   ,       **snake_case__			):

																		'''simple docstring'''


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




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


									class 					_lowercase      (bart_model.model.encoder.__class__					):
													'''simple docstring'''







													def 					_lowerCamelCase    (       self   ,       snake_case__   ,       **snake_case__			):

																		'''simple docstring'''


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




									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(__a			).numpy()
									with self.assertRaises(__a			):
														# FakeEncoder.call() accepts **kwargs -> no filtering -> value error due to unexpected input "foo"
														bart_model.generate(__a   ,       foo="bar"			)





import argparse
from collections import defaultdict

import yaml


_snake_case		=  "docs/source/en/_toctree.yml"






def 				A				(					_lowerCamelCase  ):





							'''simple docstring'''




							_lowerCAmelCase     :							Dict    =    defaultdict(_lowerCamelCase  )
							_lowerCAmelCase     :							Any    =    []
							_lowerCAmelCase     :							List[str]    =    []
							for doc in doc_list:
														if "local" in doc:
																					counts[doc["local"]] += 1

														if doc["title"].lower() == "overview":
																					overview_doc.append({"local": doc["local"], "title": doc["title"]}  )
														else:
																					new_doc_list.append(_lowerCamelCase  )

							_lowerCAmelCase     :							Optional[Any]    =    new_doc_list
							_lowerCAmelCase     :							List[Any]    =    [key for key, value in counts.items() if value > 1]

							_lowerCAmelCase     :							str    =    []
							for duplicate_key in duplicates:
														_lowerCAmelCase     :							List[str]    =    list({doc["title"] for doc in doc_list if doc["local"] == duplicate_key}  )
														if len(_lowerCamelCase  ) > 1:
																					raise ValueError(
																					    F"{duplicate_key} is present several times in the documentation table of content at "
																					    "`docs/source/en/_toctree.yml` with different *Title* values. Choose one of those and remove the "
																					    "others."  )
														# Only add this once
														new_doc.append({"local": duplicate_key, "title": titles[0]}  )

							# Add none duplicate-keys
							new_doc.extend([doc for doc in doc_list if "local" not in counts or counts[doc["local"]] == 1]  )
							_lowerCAmelCase     :							Optional[Any]    =    sorted(_lowerCamelCase				, key=lambda _lowerCamelCase  : s["title"].lower()  )

							# "overview" gets special treatment and is always first
							if len(_lowerCamelCase  ) > 1:
														raise ValueError("{doc_list} has two 'overview' docs which is not allowed."  )

							overview_doc.extend(_lowerCamelCase  )

							# Sort
							return overview_doc






def 				A				(					_lowerCamelCase=False  ):





							'''simple docstring'''




							with open(_lowerCamelCase				, encoding="utf-8"  ) as f:
														_lowerCAmelCase     :							int    =    yaml.safe_load(f.read()  )

							# Get to the API doc
							_lowerCAmelCase     :							Optional[Any]    =    0
							while content[api_idx]["title"] != "API":
														api_idx += 1
							_lowerCAmelCase     :							List[str]    =    content[api_idx]["sections"]

							# Then to the model doc
							_lowerCAmelCase     :							Union[str, Any]    =    0
							while api_doc[scheduler_idx]["title"] != "Schedulers":
														scheduler_idx += 1

							_lowerCAmelCase     :							Optional[Any]    =    api_doc[scheduler_idx]["sections"]
							_lowerCAmelCase     :							Optional[Any]    =    clean_doc_toc(_lowerCamelCase  )

							_lowerCAmelCase     :							int    =    False
							if new_scheduler_doc != scheduler_doc:
														_lowerCAmelCase     :							List[Any]    =    True
														if overwrite:
																					_lowerCAmelCase     :							Dict    =    new_scheduler_doc

							if diff:
														if overwrite:
																					_lowerCAmelCase     :							Tuple    =    api_doc
																					with open(_lowerCamelCase				, "w"				, encoding="utf-8"  ) as f:
																												f.write(yaml.dump(_lowerCamelCase				, allow_unicode=_lowerCamelCase  )  )
														else:
																					raise ValueError(
																					    "The model doc part of the table of content is not properly sorted, run `make style` to fix this."  )






def 				A				(					_lowerCamelCase=False  ):





							'''simple docstring'''




							with open(_lowerCamelCase				, encoding="utf-8"  ) as f:
														_lowerCAmelCase     :							Tuple    =    yaml.safe_load(f.read()  )

							# Get to the API doc
							_lowerCAmelCase     :							Optional[int]    =    0
							while content[api_idx]["title"] != "API":
														api_idx += 1
							_lowerCAmelCase     :							int    =    content[api_idx]["sections"]

							# Then to the model doc
							_lowerCAmelCase     :							List[str]    =    0
							while api_doc[pipeline_idx]["title"] != "Pipelines":
														pipeline_idx += 1

							_lowerCAmelCase     :							Dict    =    False
							_lowerCAmelCase     :							Optional[int]    =    api_doc[pipeline_idx]["sections"]
							_lowerCAmelCase     :							Tuple    =    []

							# sort sub pipeline docs
							for pipeline_doc in pipeline_docs:
														if "section" in pipeline_doc:
																					_lowerCAmelCase     :							List[Any]    =    pipeline_doc["section"]
																					_lowerCAmelCase     :							Union[str, Any]    =    clean_doc_toc(_lowerCamelCase  )
																					if overwrite:
																												_lowerCAmelCase     :							Optional[Any]    =    new_sub_pipeline_doc
														new_pipeline_docs.append(_lowerCamelCase  )

							# sort overall pipeline doc
							_lowerCAmelCase     :							Union[str, Any]    =    clean_doc_toc(_lowerCamelCase  )

							if new_pipeline_docs != pipeline_docs:
														_lowerCAmelCase     :							Dict    =    True
														if overwrite:
																					_lowerCAmelCase     :							Optional[int]    =    new_pipeline_docs

							if diff:
														if overwrite:
																					_lowerCAmelCase     :							Optional[int]    =    api_doc
																					with open(_lowerCamelCase				, "w"				, encoding="utf-8"  ) as f:
																												f.write(yaml.dump(_lowerCamelCase				, allow_unicode=_lowerCamelCase  )  )
														else:
																					raise ValueError(
																					    "The model doc part of the table of content is not properly sorted, run `make style` to fix this."  )


if __name__ == "__main__":
						_snake_case		=  argparse.ArgumentParser()
						parser.add_argument("--fix_and_overwrite", action="store_true", help="Whether to fix inconsistencies.")
						_snake_case		=  parser.parse_args()

						check_scheduler_doc(args.fix_and_overwrite)
						check_pipeline_doc(args.fix_and_overwrite)
