import os

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

import datasets


_lowerCAmelCase       =		datasets.logging.get_logger(__name__)


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

_lowerCAmelCase       =		"\\nBLEURT 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)\nand 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\nit for your specific application (the latter is expected to perform better).\n\nSee the project's README at https://github.com/google-research/bleurt#readme for more information.\n"

_lowerCAmelCase       =		"\nBLEURT score.\n\nArgs:\n    `predictions` (list of str): prediction/candidate sentences\n    `references` (list of str): reference sentences\n    `checkpoint` BLEURT checkpoint. Will default to BLEURT-tiny if None.\n\nReturns:\n    'scores': List of scores.\nExamples:\n\n    >>> predictions = [\"hello there\", \"general kenobi\"]\n    >>> references = [\"hello there\", \"general kenobi\"]\n    >>> bleurt = datasets.load_metric(\"bleurt\")\n    >>> results = bleurt.compute(predictions=predictions, references=references)\n    >>> print([round(v, 2) for v in results[\"scores\"]])\n    [1.03, 1.04]\n"

_lowerCAmelCase       =		{
    "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 ):







							def  UpperCamelCase_						(   self      :     Any		):
												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  UpperCamelCase_						(   self      :     List[Any] ,			_A      :     int		):
												# 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\').'''		)
																	_UpperCamelCase			     =     '''bleurt-base-128'''

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

												elif self.config_name.upper() in CHECKPOINT_URLS:
																	_UpperCamelCase			     =     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
												_UpperCamelCase			     =     dl_manager.download_and_extract(CHECKPOINT_URLS[checkpoint_name]		)
												_UpperCamelCase			     =     score.BleurtScorer(os.path.join(_A ,			_A		)		)







							def  UpperCamelCase_						(   self      :     Dict ,			_A      :     List[str] ,			_A      :     Dict		):
												_UpperCamelCase			     =     self.scorer.score(references=_A ,			candidates=_A		)
												return {"scores": scores}

import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
class     lowerCAmelCase_     (       __lowercase,  unittest.TestCase ):
							UpperCAmelCase				   =     ShapEPipeline
							UpperCAmelCase				   =     ["prompt"]
							UpperCAmelCase				   =     ["prompt"]
							UpperCAmelCase				   =     [
							    "num_images_per_prompt",
							    "num_inference_steps",
							    "generator",
							    "latents",
							    "guidance_scale",
							    "frame_size",
							    "output_type",
							    "return_dict",
							]
							UpperCAmelCase				   =     False







							@property
							def  UpperCamelCase_						(   self      :     Union[str, Any]		):
												return 32







							@property
							def  UpperCamelCase_						(   self      :     int		):
												return 32







							@property
							def  UpperCamelCase_						(   self      :     List[str]		):
												return self.time_input_dim * 4







							@property
							def  UpperCamelCase_						(   self      :     Optional[Any]		):
												return 8







							@property
							def  UpperCamelCase_						(   self      :     int		):
												_UpperCamelCase			     =     CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)
												return tokenizer







							@property
							def  UpperCamelCase_						(   self      :     List[Any]		):
												torch.manual_seed(0		)
												_UpperCamelCase			     =     CLIPTextConfig(
												    bos_token_id=0 ,			eos_token_id=2 ,			hidden_size=self.text_embedder_hidden_size ,			projection_dim=self.text_embedder_hidden_size ,			intermediate_size=37 ,			layer_norm_eps=1e-05 ,			num_attention_heads=4 ,			num_hidden_layers=5 ,			pad_token_id=1 ,			vocab_size=1000 ,			)
												return CLIPTextModelWithProjection(_A		)







							@property
							def  UpperCamelCase_						(   self      :     int		):
												torch.manual_seed(0		)

												_UpperCamelCase			     =     {
												    '''num_attention_heads''': 2,
												    '''attention_head_dim''': 16,
												    '''embedding_dim''': self.time_input_dim,
												    '''num_embeddings''': 32,
												    '''embedding_proj_dim''': self.text_embedder_hidden_size,
												    '''time_embed_dim''': self.time_embed_dim,
												    '''num_layers''': 1,
												    '''clip_embed_dim''': self.time_input_dim * 2,
												    '''additional_embeddings''': 0,
												    '''time_embed_act_fn''': '''gelu''',
												    '''norm_in_type''': '''layer''',
												    '''encoder_hid_proj_type''': None,
												    '''added_emb_type''': None,
												}

												_UpperCamelCase			     =     PriorTransformer(**_A		)
												return model







							@property
							def  UpperCamelCase_						(   self      :     Union[str, Any]		):
												torch.manual_seed(0		)

												_UpperCamelCase			     =     {
												    '''param_shapes''': (
												        (self.renderer_dim, 93),
												        (self.renderer_dim, 8),
												        (self.renderer_dim, 8),
												        (self.renderer_dim, 8),
												    ),
												    '''d_latent''': self.time_input_dim,
												    '''d_hidden''': self.renderer_dim,
												    '''n_output''': 12,
												    '''background''': (
												        0.1,
												        0.1,
												        0.1,
												    ),
												}
												_UpperCamelCase			     =     ShapERenderer(**_A		)
												return model







							def  UpperCamelCase_						(   self      :     str		):
												_UpperCamelCase			     =     self.dummy_prior
												_UpperCamelCase			     =     self.dummy_text_encoder
												_UpperCamelCase			     =     self.dummy_tokenizer
												_UpperCamelCase			     =     self.dummy_renderer

												_UpperCamelCase			     =     HeunDiscreteScheduler(
												    beta_schedule='''exp''' ,			num_train_timesteps=1024 ,			prediction_type='''sample''' ,			use_karras_sigmas=_A ,			clip_sample=_A ,			clip_sample_range=1.0 ,			)
												_UpperCamelCase			     =     {
												    '''prior''': prior,
												    '''text_encoder''': text_encoder,
												    '''tokenizer''': tokenizer,
												    '''renderer''': renderer,
												    '''scheduler''': scheduler,
												}

												return components







							def  UpperCamelCase_						(   self      :     Tuple ,			_A      :     Tuple ,			_A      :     Optional[int]=0		):
												if str(_A		).startswith('''mps'''		):
																	_UpperCamelCase			     =     torch.manual_seed(_A		)
												else:
																	_UpperCamelCase			     =     torch.Generator(device=_A		).manual_seed(_A		)
												_UpperCamelCase			     =     {
												    '''prompt''': '''horse''',
												    '''generator''': generator,
												    '''num_inference_steps''': 1,
												    '''frame_size''': 32,
												    '''output_type''': '''np''',
												}
												return inputs







							def  UpperCamelCase_						(   self      :     Any		):
												_UpperCamelCase			     =     '''cpu'''

												_UpperCamelCase			     =     self.get_dummy_components()

												_UpperCamelCase			     =     self.pipeline_class(**_A		)
												_UpperCamelCase			     =     pipe.to(_A		)

												pipe.set_progress_bar_config(disable=_A		)

												_UpperCamelCase			     =     pipe(**self.get_dummy_inputs(_A		)		)
												_UpperCamelCase			     =     output.images[0]
												_UpperCamelCase			     =     image[0, -3:, -3:, -1]

												assert image.shape == (20, 32, 32, 3)

												_UpperCamelCase			     =     np.array(
												    [
												        0.0003_9216,
												        0.0003_9216,
												        0.0003_9216,
												        0.0003_9216,
												        0.0003_9216,
												        0.0003_9216,
												        0.0003_9216,
												        0.0003_9216,
												        0.0003_9216,
												    ]		)

												assert np.abs(image_slice.flatten() - expected_slice		).max() < 1e-2







							def  UpperCamelCase_						(   self      :     Any		):
												# NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
												self._test_inference_batch_consistent(batch_sizes=[1, 2]		)







							def  UpperCamelCase_						(   self      :     Any		):
												_UpperCamelCase			     =     torch_device == '''cpu'''
												_UpperCamelCase			     =     True

												self._test_inference_batch_single_identical(
												    batch_size=2 ,			test_max_difference=_A ,			relax_max_difference=_A ,			)







							def  UpperCamelCase_						(   self      :     Any		):
												_UpperCamelCase			     =     self.get_dummy_components()
												_UpperCamelCase			     =     self.pipeline_class(**_A		)
												_UpperCamelCase			     =     pipe.to(_A		)
												pipe.set_progress_bar_config(disable=_A		)

												_UpperCamelCase			     =     1
												_UpperCamelCase			     =     2

												_UpperCamelCase			     =     self.get_dummy_inputs(_A		)

												for key in inputs.keys():
																	if key in self.batch_params:
																						_UpperCamelCase			     =     batch_size * [inputs[key]]

												_UpperCamelCase			     =     pipe(**_A ,			num_images_per_prompt=_A		)[0]

												assert images.shape[0] == batch_size * num_images_per_prompt




@slow
@require_torch_gpu
class     lowerCAmelCase_     (       unittest.TestCase ):







							def  UpperCamelCase_						(   self      :     str		):
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()







							def  UpperCamelCase_						(   self      :     List[str]		):
												_UpperCamelCase			     =     load_numpy(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
												    '''/shap_e/test_shap_e_np_out.npy'''		)
												_UpperCamelCase			     =     ShapEPipeline.from_pretrained('''openai/shap-e'''		)
												_UpperCamelCase			     =     pipe.to(_A		)
												pipe.set_progress_bar_config(disable=_A		)

												_UpperCamelCase			     =     torch.Generator(device=_A		).manual_seed(0		)

												_UpperCamelCase			     =     pipe(
												    '''a shark''' ,			generator=_A ,			guidance_scale=15.0 ,			num_inference_steps=64 ,			frame_size=64 ,			output_type='''np''' ,			).images[0]

												assert images.shape == (20, 64, 64, 3)

												assert_mean_pixel_difference(_A ,			_A		)
