





import os

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

import datasets


__UpperCamelCase:	Dict = datasets.logging.get_logger(__name__)


__UpperCamelCase:	str = """\
@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}
}
"""

__UpperCamelCase:	Dict = """\
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.
"""

__UpperCamelCase:	Optional[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]
"""

__UpperCamelCase:	Optional[int] = {
    """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    snake_case__(						self:						str			):
			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    snake_case__(						self:						Optional[Any],					lowerCamelCase_:						Union[str, Any]			):
			# check that config name specifies a valid BLEURT model
			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__					:      Tuple      =		'bleurt-base-128'

			if self.config_name.lower() in CHECKPOINT_URLS:
					lowercase__					:      Optional[int]      =		self.config_name.lower()

			elif self.config_name.upper() in CHECKPOINT_URLS:
					lowercase__					:      str      =		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__					:      List[str]      =		dl_manager.download_and_extract(CHECKPOINT_URLS[checkpoint_name]			)
			lowercase__					:      Optional[int]      =		score.BleurtScorer(os.path.join(lowerCamelCase_,					lowerCamelCase_			)			)





	def    snake_case__(						self:						Union[str, Any],					lowerCamelCase_:						List[str],					lowerCamelCase_:						int			):
			lowercase__					:      List[Any]      =		self.scorer.score(references=lowerCamelCase_,					candidates=lowerCamelCase_			)
			return {"scores": scores}









import unittest

from parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
			import torch

			from transformers import OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel


class       __lowerCAmelCase					:


	'''simple docstring'''





	def __init__(						self:						Any,					lowerCamelCase_:						Optional[Any],					lowerCamelCase_:						List[str]=13,					lowerCamelCase_:						Optional[Any]=7,					lowerCamelCase_:						Optional[Any]=True,					lowerCamelCase_:						Tuple=True,					lowerCamelCase_:						Any=False,					lowerCamelCase_:						Union[str, Any]=True,					lowerCamelCase_:						Optional[Any]=99,					lowerCamelCase_:						Tuple=32,					lowerCamelCase_:						Any=5,					lowerCamelCase_:						Tuple=4,					lowerCamelCase_:						List[Any]=37,					lowerCamelCase_:						Union[str, Any]="gelu",					lowerCamelCase_:						str=0.1,					lowerCamelCase_:						Union[str, Any]=0.1,					lowerCamelCase_:						Any=512,					lowerCamelCase_:						Union[str, Any]=16,					lowerCamelCase_:						Any=2,					lowerCamelCase_:						str=0.0_2,					lowerCamelCase_:						Union[str, Any]=3,					lowerCamelCase_:						List[str]=4,					lowerCamelCase_:						Tuple=None,					):
			lowercase__					:      List[str]      =		parent
			lowercase__					:      str      =		batch_size
			lowercase__					:      Optional[int]      =		seq_length
			lowercase__					:      List[Any]      =		is_training
			lowercase__					:      List[str]      =		use_input_mask
			lowercase__					:      Optional[int]      =		use_token_type_ids
			lowercase__					:      str      =		use_labels
			lowercase__					:      Optional[Any]      =		vocab_size
			lowercase__					:      Dict      =		hidden_size
			lowercase__					:      List[Any]      =		num_hidden_layers
			lowercase__					:      Tuple      =		num_attention_heads
			lowercase__					:      Dict      =		intermediate_size
			lowercase__					:      Tuple      =		hidden_act
			lowercase__					:      List[str]      =		hidden_dropout_prob
			lowercase__					:      List[str]      =		attention_probs_dropout_prob
			lowercase__					:      str      =		max_position_embeddings
			lowercase__					:      List[Any]      =		type_vocab_size
			lowercase__					:      Tuple      =		type_sequence_label_size
			lowercase__					:      Optional[int]      =		initializer_range
			lowercase__					:      str      =		num_labels
			lowercase__					:      Optional[int]      =		num_choices
			lowercase__					:      Dict      =		scope

	def    snake_case__(						self:						Union[str, Any]			):
			lowercase__					:      Any      =		ids_tensor([self.batch_size, self.seq_length],					self.vocab_size			)

			lowercase__					:      List[Any]      =		None
			if self.use_input_mask:
					lowercase__					:      List[str]      =		random_attention_mask([self.batch_size, self.seq_length]			)

			lowercase__					:      Optional[Any]      =		None
			if self.use_token_type_ids:
					lowercase__					:      Tuple      =		ids_tensor([self.batch_size, self.seq_length],					self.type_vocab_size			)

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

			lowercase__					:      Tuple      =		self.get_config()

			return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

	def    snake_case__(						self:						Tuple			):
			return OpenLlamaConfig(
			    vocab_size=self.vocab_size,					hidden_size=self.hidden_size,					num_hidden_layers=self.num_hidden_layers,					num_attention_heads=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,					max_position_embeddings=self.max_position_embeddings,					type_vocab_size=self.type_vocab_size,					is_decoder=lowerCamelCase_,					initializer_range=self.initializer_range,					use_stable_embedding=lowerCamelCase_,					)

	def    snake_case__(						self:						Optional[int],					lowerCamelCase_:						Optional[int],					lowerCamelCase_:						List[Any],					lowerCamelCase_:						List[str],					lowerCamelCase_:						Optional[int],					lowerCamelCase_:						Dict,					lowerCamelCase_:						Optional[int],					lowerCamelCase_:						str			):
			lowercase__					:      Union[str, Any]      =		OpenLlamaModel(config=lowerCamelCase_			)
			model.to(lowerCamelCase_			)
			model.eval()
			lowercase__					:      Union[str, Any]      =		model(lowerCamelCase_,					attention_mask=lowerCamelCase_			)
			lowercase__					:      List[Any]      =		model(lowerCamelCase_			)
			self.parent.assertEqual(result.last_hidden_state.shape,					(self.batch_size, self.seq_length, self.hidden_size)			)

	def    snake_case__(						self:						str,					lowerCamelCase_:						Optional[Any],					lowerCamelCase_:						str,					lowerCamelCase_:						str,					lowerCamelCase_:						Optional[Any],					lowerCamelCase_:						Tuple,					lowerCamelCase_:						Optional[Any],					lowerCamelCase_:						List[str],					lowerCamelCase_:						Dict,					lowerCamelCase_:						Union[str, Any],					):
			lowercase__					:      Tuple      =		True
			lowercase__					:      int      =		OpenLlamaModel(lowerCamelCase_			)
			model.to(lowerCamelCase_			)
			model.eval()
			lowercase__					:      List[Any]      =		model(
			    lowerCamelCase_,					attention_mask=lowerCamelCase_,					encoder_hidden_states=lowerCamelCase_,					encoder_attention_mask=lowerCamelCase_,					)
			lowercase__					:      Optional[int]      =		model(
			    lowerCamelCase_,					attention_mask=lowerCamelCase_,					encoder_hidden_states=lowerCamelCase_,					)
			lowercase__					:      str      =		model(lowerCamelCase_,					attention_mask=lowerCamelCase_			)
			self.parent.assertEqual(result.last_hidden_state.shape,					(self.batch_size, self.seq_length, self.hidden_size)			)

	def    snake_case__(						self:						List[Any],					lowerCamelCase_:						Optional[Any],					lowerCamelCase_:						Any,					lowerCamelCase_:						str,					lowerCamelCase_:						List[str],					lowerCamelCase_:						Any,					lowerCamelCase_:						Dict,					lowerCamelCase_:						int,					lowerCamelCase_:						Any,					lowerCamelCase_:						str,					):
			lowercase__					:      Optional[Any]      =		OpenLlamaForCausalLM(config=lowerCamelCase_			)
			model.to(lowerCamelCase_			)
			model.eval()
			lowercase__					:      List[Any]      =		model(lowerCamelCase_,					attention_mask=lowerCamelCase_,					labels=lowerCamelCase_			)
			self.parent.assertEqual(result.logits.shape,					(self.batch_size, self.seq_length, self.vocab_size)			)

	def    snake_case__(						self:						Tuple,					lowerCamelCase_:						str,					lowerCamelCase_:						Optional[Any],					lowerCamelCase_:						Optional[int],					lowerCamelCase_:						str,					lowerCamelCase_:						List[Any],					lowerCamelCase_:						List[Any],					lowerCamelCase_:						str,					lowerCamelCase_:						Optional[Any],					lowerCamelCase_:						Union[str, Any],					):
			lowercase__					:      Optional[int]      =		True
			lowercase__					:      Optional[int]      =		True
			lowercase__					:      Tuple      =		OpenLlamaForCausalLM(config=lowerCamelCase_			)
			model.to(lowerCamelCase_			)
			model.eval()

			# first forward pass
			lowercase__					:      Any      =		model(
			    lowerCamelCase_,					attention_mask=lowerCamelCase_,					encoder_hidden_states=lowerCamelCase_,					encoder_attention_mask=lowerCamelCase_,					use_cache=lowerCamelCase_,					)
			lowercase__					:      str      =		outputs.past_key_values

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

			# append to next input_ids and
			lowercase__					:      Tuple      =		torch.cat([input_ids, next_tokens],					dim=-1			)
			lowercase__					:      List[str]      =		torch.cat([input_mask, next_mask],					dim=-1			)

			lowercase__					:      int      =		model(
			    lowerCamelCase_,					attention_mask=lowerCamelCase_,					encoder_hidden_states=lowerCamelCase_,					encoder_attention_mask=lowerCamelCase_,					output_hidden_states=lowerCamelCase_,					)['hidden_states'][0]
			lowercase__					:      int      =		model(
			    lowerCamelCase_,					attention_mask=lowerCamelCase_,					encoder_hidden_states=lowerCamelCase_,					encoder_attention_mask=lowerCamelCase_,					past_key_values=lowerCamelCase_,					output_hidden_states=lowerCamelCase_,					)['hidden_states'][0]

			# select random slice
			lowercase__					:      Tuple      =		ids_tensor((1,),					output_from_past.shape[-1]			).item()
			lowercase__					:      Union[str, Any]      =		output_from_no_past[:, -3:, random_slice_idx].detach()
			lowercase__					:      Tuple      =		output_from_past[:, :, random_slice_idx].detach()

			self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]			)

			# test that outputs are equal for slice
			self.parent.assertTrue(torch.allclose(lowerCamelCase_,					lowerCamelCase_,					atol=1E-3			)			)





	def    snake_case__(						self:						Optional[Any]			):
			lowercase__					:      Dict      =		self.prepare_config_and_inputs()
			(
			    (
			    lowercase__
			)   ,				(
			    lowercase__
			)   ,				(
			    lowercase__
			)   ,				(
			    lowercase__
			)   ,				(
			    lowercase__
			)   ,				(
			    lowercase__
			)   ,				(
			    lowercase__
			)   ,				
			)					:      List[str]      =		config_and_inputs
			lowercase__					:      Any      =		{'input_ids': input_ids, 'attention_mask': input_mask}
			return config, inputs_dict


@require_torch
class       __lowerCAmelCase					(       _UpperCamelCase   ,    _UpperCamelCase   ,    _UpperCamelCase   ,    unittest.TestCase      ):


	'''simple docstring'''





	_A													=				(
	    (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
	)
	_A													=				(OpenLlamaForCausalLM,) if is_torch_available() else ()
	_A													=				(
	    {
	        "feature-extraction": OpenLlamaModel,
	        "text-classification": OpenLlamaForSequenceClassification,
	        "text-generation": OpenLlamaForCausalLM,
	        "zero-shot": OpenLlamaForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	_A													=				False
	_A													=				False

	def    snake_case__(						self:						Any			):
			lowercase__					:      List[Any]      =		OpenLlamaModelTester(self			)
			lowercase__					:      Dict      =		ConfigTester(self,					config_class=lowerCamelCase_,					hidden_size=37			)

	def    snake_case__(						self:						Any			):
			self.config_tester.run_common_tests()

	def    snake_case__(						self:						Dict			):
			lowercase__					:      Any      =		self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*lowerCamelCase_			)

	def    snake_case__(						self:						str			):
			lowercase__					:      Any      =		self.model_tester.prepare_config_and_inputs()
			for type in ["absolute", "relative_key", "relative_key_query"]:
					lowercase__					:      Dict      =		type
					self.model_tester.create_and_check_model(*lowerCamelCase_			)

	def    snake_case__(						self:						Any			):
			lowercase__   ,				lowercase__					:      List[Any]      =		self.model_tester.prepare_config_and_inputs_for_common()
			lowercase__					:      Tuple      =		3
			lowercase__					:      Union[str, Any]      =		input_dict['input_ids']
			lowercase__					:      Union[str, Any]      =		input_ids.ne(1			).to(lowerCamelCase_			)
			lowercase__					:      int      =		ids_tensor([self.model_tester.batch_size],					self.model_tester.type_sequence_label_size			)
			lowercase__					:      Optional[Any]      =		OpenLlamaForSequenceClassification(lowerCamelCase_			)
			model.to(lowerCamelCase_			)
			model.eval()
			lowercase__					:      Optional[Any]      =		model(lowerCamelCase_,					attention_mask=lowerCamelCase_,					labels=lowerCamelCase_			)
			self.assertEqual(result.logits.shape,					(self.model_tester.batch_size, self.model_tester.num_labels)			)

	def    snake_case__(						self:						Any			):
			lowercase__   ,				lowercase__					:      int      =		self.model_tester.prepare_config_and_inputs_for_common()
			lowercase__					:      str      =		3
			lowercase__					:      Optional[Any]      =		'single_label_classification'
			lowercase__					:      Union[str, Any]      =		input_dict['input_ids']
			lowercase__					:      str      =		input_ids.ne(1			).to(lowerCamelCase_			)
			lowercase__					:      Any      =		ids_tensor([self.model_tester.batch_size],					self.model_tester.type_sequence_label_size			)
			lowercase__					:      Dict      =		OpenLlamaForSequenceClassification(lowerCamelCase_			)
			model.to(lowerCamelCase_			)
			model.eval()
			lowercase__					:      Optional[int]      =		model(lowerCamelCase_,					attention_mask=lowerCamelCase_,					labels=lowerCamelCase_			)
			self.assertEqual(result.logits.shape,					(self.model_tester.batch_size, self.model_tester.num_labels)			)

	def    snake_case__(						self:						Union[str, Any]			):
			lowercase__   ,				lowercase__					:      Optional[int]      =		self.model_tester.prepare_config_and_inputs_for_common()
			lowercase__					:      Dict      =		3
			lowercase__					:      List[Any]      =		'multi_label_classification'
			lowercase__					:      Any      =		input_dict['input_ids']
			lowercase__					:      int      =		input_ids.ne(1			).to(lowerCamelCase_			)
			lowercase__					:      Optional[int]      =		ids_tensor(
			    [self.model_tester.batch_size, config.num_labels],					self.model_tester.type_sequence_label_size			).to(torch.float			)
			lowercase__					:      Optional[int]      =		OpenLlamaForSequenceClassification(lowerCamelCase_			)
			model.to(lowerCamelCase_			)
			model.eval()
			lowercase__					:      List[Any]      =		model(lowerCamelCase_,					attention_mask=lowerCamelCase_,					labels=lowerCamelCase_			)
			self.assertEqual(result.logits.shape,					(self.model_tester.batch_size, self.model_tester.num_labels)			)

	@unittest.skip('Open-Llama buffers include complex numbers, which breaks this test'			)
	def    snake_case__(						self:						List[str]			):
			pass





	@parameterized.expand([('linear',), ('dynamic',)]			)
	def    snake_case__(						self:						Any,					lowerCamelCase_:						List[Any]			):
			lowercase__   ,				lowercase__					:      Optional[Any]      =		self.model_tester.prepare_config_and_inputs_for_common()
			lowercase__					:      Optional[int]      =		ids_tensor([1, 10],					config.vocab_size			)
			lowercase__					:      Optional[Any]      =		ids_tensor([1, int(config.max_position_embeddings * 1.5			)],					config.vocab_size			)

			set_seed(42			)  # Fixed seed at init time so the two models get the same random weights
			lowercase__					:      str      =		OpenLlamaModel(lowerCamelCase_			)
			original_model.to(lowerCamelCase_			)
			original_model.eval()
			lowercase__					:      Union[str, Any]      =		original_model(lowerCamelCase_			).last_hidden_state
			lowercase__					:      List[str]      =		original_model(lowerCamelCase_			).last_hidden_state

			set_seed(42			)  # Fixed seed at init time so the two models get the same random weights
			lowercase__					:      List[str]      =		{'type': scaling_type, 'factor': 1_0.0}
			lowercase__					:      Union[str, Any]      =		OpenLlamaModel(lowerCamelCase_			)
			scaled_model.to(lowerCamelCase_			)
			scaled_model.eval()
			lowercase__					:      Dict      =		scaled_model(lowerCamelCase_			).last_hidden_state
			lowercase__					:      Union[str, Any]      =		scaled_model(lowerCamelCase_			).last_hidden_state

			# Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
			# maximum sequence length, so the outputs for the short input should match.
			if scaling_type == "dynamic":
					self.assertTrue(torch.allclose(lowerCamelCase_,					lowerCamelCase_,					atol=1E-5			)			)
			else:
					self.assertFalse(torch.allclose(lowerCamelCase_,					lowerCamelCase_,					atol=1E-5			)			)

			# The output should be different for long inputs
			self.assertFalse(torch.allclose(lowerCamelCase_,					lowerCamelCase_,					atol=1E-5			)			)


