


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



					'''simple docstring'''







					def A      (				self						)   ->			List[Any]:







								'''simple docstring'''

								__lowercase 					=   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 ,       )

								__lowercase 					=   tf.convert_to_tensor(
								    [[0, 0], [0, 9], [0, 1_0], [0, 2_5], [0, 2_6], [1, 1_3], [1, 1_7], [1, 1_8], [1, 2_0], [1, 2_7]] ,       dtype=tf.intaa ,       )  # expected non filtered idx as noted above

								__lowercase 					=   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

								__lowercase 					=   tf_top_k_top_p_filtering(snake_case_ ,       top_k=1_0 ,       top_p=0.6 ,       min_tokens_to_keep=4						)

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

								tf.debugging.assert_near(snake_case_ ,       snake_case_ ,       rtol=1e-1_2						)
								tf.debugging.assert_equal(snake_case_ ,       snake_case_						)




@require_tf
class 	lowerCamelCase_       (				unittest.TestCase     ,     lowerCAmelCase__				):



					'''simple docstring'''







					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 A      (				self						)   ->			Optional[Any]:







								'''simple docstring'''

								__lowercase 					=   TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''						)
								__lowercase 					=   2
								__lowercase 					=   2
								class 	lowerCamelCase_       (				tf.Module				):



													'''simple docstring'''







													def __init__(				self ,       snake_case_						)   ->			str:







																'''simple docstring'''

																super(snake_case_ ,       self						).__init__()
																__lowercase 					=   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=snake_case_ ,       )
													def A      (				self ,       snake_case_ ,       snake_case_						)   ->			str:







																'''simple docstring'''

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




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






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







								'''simple docstring'''

								__lowercase 					=   TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''						)
								__lowercase 					=   1
								__lowercase 					=   2
								class 	lowerCamelCase_       (				tf.Module				):



													'''simple docstring'''







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







																'''simple docstring'''

																super(snake_case_ ,       self						).__init__()
																__lowercase 					=   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=snake_case_ ,       )
													def A      (				self ,       snake_case_ ,       snake_case_						)   ->			Optional[Any]:







																'''simple docstring'''

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




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






					@slow
					@require_tensorflow_text
					def A      (				self						)   ->			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=snake_case_						)
											class 	lowerCamelCase_       (				tf.keras.layers.Layer				):



																'''simple docstring'''







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







																			'''simple docstring'''

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






																def A      (				self ,       snake_case_ ,       *snake_case_ ,       **snake_case_						)   ->			List[Any]:







																			'''simple docstring'''

																			__lowercase 					=   self.tokenizer.tokenize(snake_case_						)
																			__lowercase					, __lowercase 					=   text.pad_model_inputs(
																			    snake_case_ ,       max_seq_length=6_4 ,       pad_value=self.model.config.pad_token_id						)
																			__lowercase 					=   self.model.generate(input_ids=snake_case_ ,       attention_mask=snake_case_						)
																			return self.tokenizer.detokenize(snake_case_						)




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






					def A      (				self						)   ->			List[Any]:







								'''simple docstring'''

								__lowercase 					=   {
								    '''do_sample''': True,
								    '''num_beams''': 1,
								    '''top_p''': 0.7,
								    '''top_k''': 1_0,
								    '''temperature''': 0.7,
								}
								__lowercase 					=   1_4

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

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

								__lowercase 					=   [6_3_8, 1_9_8]
								with tf.device(''':/CPU:0'''						):
											tf.random.set_seed(0						)
											__lowercase 					=   model.generate(**snake_case_ ,       eos_token_id=snake_case_ ,       **snake_case_						)
								self.assertTrue(expectation == len(generated_tokens[0]						)						)






					def A      (				self						)   ->			Optional[Any]:







								'''simple docstring'''

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



													'''simple docstring'''







													def A      (				self ,       snake_case_ ,       snake_case_=None ,       **snake_case_						)   ->			Optional[int]:







																'''simple docstring'''

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




								__lowercase 					=   FakeBart.from_pretrained('''hf-internal-testing/tiny-random-bart'''						)
								__lowercase 					=   bart_model.generate(snake_case_ ,       foo='''bar'''						).numpy()
								self.assertTrue(np.array_equal(snake_case_ ,       snake_case_						)						)
								class 	lowerCamelCase_       (				bart_model.model.encoder.__class__				):



													'''simple docstring'''







													def A      (				self ,       snake_case_ ,       **snake_case_						)   ->			Any:







																'''simple docstring'''

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




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

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






import os

from bleurt import score  # From: git+https://github.com/google-research/bleurt.git

import datasets


a : Optional[int]		=				datasets.logging.get_logger(__name__)


a : Tuple		=				'''\
@inproceedings{bleurt,
  title={BLEURT: Learning Robust Metrics for Text Generation},
  author={Thibault Sellam and Dipanjan Das and Ankur P. Parikh},
  booktitle={ACL},
  year={2020},
  url={https://arxiv.org/abs/2004.04696}
}
'''

a : Union[str, Any]		=				'''\
BLEURT a learnt evaluation metric for Natural Language Generation. It is built using multiple phases of transfer learning starting from a pretrained BERT model (Devlin et al. 2018)
and then employing another pre-training phrase using synthetic data. Finally it is trained on WMT human annotations. You may run BLEURT out-of-the-box or fine-tune
it for your specific application (the latter is expected to perform better).

See the project\'s README at https://github.com/google-research/bleurt#readme for more information.
'''

a : Union[str, Any]		=				'''
BLEURT score.

Args:
    `predictions` (list of str): prediction/candidate sentences
    `references` (list of str): reference sentences
    `checkpoint` BLEURT checkpoint. Will default to BLEURT-tiny if None.

Returns:
    \'scores\': List of scores.
Examples:

    >>> predictions = ["hello there", "general kenobi"]
    >>> references = ["hello there", "general kenobi"]
    >>> bleurt = datasets.load_metric("bleurt")
    >>> results = bleurt.compute(predictions=predictions, references=references)
    >>> print([round(v, 2) for v in results["scores"]])
    [1.03, 1.04]
'''

a : Tuple		=				{
    '''bleurt-tiny-128''': '''https://storage.googleapis.com/bleurt-oss/bleurt-tiny-128.zip''',
    '''bleurt-tiny-512''': '''https://storage.googleapis.com/bleurt-oss/bleurt-tiny-512.zip''',
    '''bleurt-base-128''': '''https://storage.googleapis.com/bleurt-oss/bleurt-base-128.zip''',
    '''bleurt-base-512''': '''https://storage.googleapis.com/bleurt-oss/bleurt-base-512.zip''',
    '''bleurt-large-128''': '''https://storage.googleapis.com/bleurt-oss/bleurt-large-128.zip''',
    '''bleurt-large-512''': '''https://storage.googleapis.com/bleurt-oss/bleurt-large-512.zip''',
    '''BLEURT-20-D3''': '''https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D3.zip''',
    '''BLEURT-20-D6''': '''https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D6.zip''',
    '''BLEURT-20-D12''': '''https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D12.zip''',
    '''BLEURT-20''': '''https://storage.googleapis.com/bleurt-oss-21/BLEURT-20.zip''',
}
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION     ,     _KWARGS_DESCRIPTION				)
class 	lowerCamelCase_       (				datasets.Metric				):



					'''simple docstring'''







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







								'''simple docstring'''

								return datasets.MetricInfo(
								    description=_DESCRIPTION ,       citation=_CITATION ,       homepage='''https://github.com/google-research/bleurt''' ,       inputs_description=_KWARGS_DESCRIPTION ,       features=datasets.Features(
								        {
								            '''predictions''': datasets.Value('''string''' ,       id='''sequence'''						),
								            '''references''': datasets.Value('''string''' ,       id='''sequence'''						),
								        }						) ,       codebase_urls=['''https://github.com/google-research/bleurt'''] ,       reference_urls=['''https://github.com/google-research/bleurt''', '''https://arxiv.org/abs/2004.04696'''] ,       )






					def A      (				self ,       snake_case_						)   ->			List[str]:







								'''simple docstring'''

								if self.config_name == "default":
											logger.warning(
											    '''Using default BLEURT-Base checkpoint for sequence maximum length 128. '''
											    '''You can use a bigger model for better results with e.g.: datasets.load_metric(\'bleurt\', \'bleurt-large-512\').'''						)
											__lowercase 					=   '''bleurt-base-128'''

								if self.config_name.lower() in CHECKPOINT_URLS:
											__lowercase 					=   self.config_name.lower()

								elif self.config_name.upper() in CHECKPOINT_URLS:
											__lowercase 					=   self.config_name.upper()

								else:
											raise KeyError(
											    F'{self.config_name} model not found. You should supply the name of a model checkpoint for bleurt in {CHECKPOINT_URLS.keys()}'						)

								# download the model checkpoint specified by self.config_name and set up the scorer
								__lowercase 					=   dl_manager.download_and_extract(CHECKPOINT_URLS[checkpoint_name]						)
								__lowercase 					=   score.BleurtScorer(os.path.join(snake_case_ ,       snake_case_						)						)






					def A      (				self ,       snake_case_ ,       snake_case_						)   ->			Tuple:







								'''simple docstring'''

								__lowercase 					=   self.scorer.score(references=snake_case_ ,       candidates=snake_case_						)
								return {"scores": scores}


