





from dataclasses import dataclass
from typing import Dict, Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .attention_processor import AttentionProcessor, AttnProcessor
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder
@dataclass
class      UpperCamelCase_						(					_lowerCamelCase   ):
					lowerCAmelCase_							=		42





class      UpperCamelCase_						(					_lowerCamelCase							,					_lowerCamelCase   ):
					lowerCAmelCase_							=		True




					@register_to_config
					def __init__(						self			,  lowerCAmelCase_ = 3			,  lowerCAmelCase_ = 3			,  lowerCAmelCase_ = ("DownEncoderBlock2D",)			,  lowerCAmelCase_ = ("UpDecoderBlock2D",)			,  lowerCAmelCase_ = (64,)			,  lowerCAmelCase_ = 1			,  lowerCAmelCase_ = "silu"			,  lowerCAmelCase_ = 4			,  lowerCAmelCase_ = 32			,  lowerCAmelCase_ = 32			,  lowerCAmelCase_ = 0.1_82_15			,  )						->							Any:
										super().__init__()

										# pass init params to Encoder
										_snake_case													=			Encoder(
										    in_channels=lowerCAmelCase_			,  out_channels=lowerCAmelCase_			,  down_block_types=lowerCAmelCase_			,  block_out_channels=lowerCAmelCase_			,  layers_per_block=lowerCAmelCase_			,  act_fn=lowerCAmelCase_			,  norm_num_groups=lowerCAmelCase_			,  double_z=lowerCAmelCase_			,  )

										# pass init params to Decoder
										_snake_case													=			Decoder(
										    in_channels=lowerCAmelCase_			,  out_channels=lowerCAmelCase_			,  up_block_types=lowerCAmelCase_			,  block_out_channels=lowerCAmelCase_			,  layers_per_block=lowerCAmelCase_			,  norm_num_groups=lowerCAmelCase_			,  act_fn=lowerCAmelCase_			,  )

										_snake_case													=			nn.Convad(2 * latent_channels			,  2 * latent_channels			,  1					)
										_snake_case													=			nn.Convad(lowerCAmelCase_			,  lowerCAmelCase_			,  1					)

										_snake_case													=			False
										_snake_case													=			False

										# only relevant if vae tiling is enabled
										_snake_case													=			self.config.sample_size
										_snake_case													=			(
										    self.config.sample_size[0]
										    if isinstance(self.config.sample_size			,  (list, tuple)					)
										    else self.config.sample_size
										)
										_snake_case													=			int(sample_size / (2 ** (len(self.config.block_out_channels					) - 1))					)
										_snake_case													=			0.25




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_=False					)						->							Optional[int]:
										if isinstance(lowerCAmelCase_			,  (Encoder, Decoder)					):
															_snake_case													=			value




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_ = True					)						->							List[Any]:
										_snake_case													=			use_tiling




					def 		lowerCAmelCase     (						self					)						->							Optional[Any]:
										self.enable_tiling(lowerCAmelCase_					)




					def 		lowerCAmelCase     (						self					)						->							List[Any]:
										_snake_case													=			True




					def 		lowerCAmelCase     (						self					)						->							Any:
										_snake_case													=			False




					@property
					# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
					def 		lowerCAmelCase     (						self					)						->							Dict[str, AttentionProcessor]:
										_snake_case													=			{}

										def fn_recursive_add_processors(lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					):
															if hasattr(lowerCAmelCase_			,  'set_processor'					):
																				_snake_case													=			module.processor

															for sub_name, child in module.named_children():
																				fn_recursive_add_processors(F'''{name}.{sub_name}'''			,  lowerCAmelCase_			,  lowerCAmelCase_					)

															return processors

										for name, module in self.named_children():
															fn_recursive_add_processors(lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)

										return processors




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_					)						->							Tuple:
										_snake_case													=			len(self.attn_processors.keys()					)

										if isinstance(lowerCAmelCase_			,  lowerCAmelCase_					) and len(lowerCAmelCase_					) != count:
															raise ValueError(
															    F'''A dict of processors was passed, but the number of processors {len(lowerCAmelCase_					)} does not match the'''
															    F''' number of attention layers: {count}. Please make sure to pass {count} processor classes.'''					)

										def fn_recursive_attn_processor(lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					):
															if hasattr(lowerCAmelCase_			,  'set_processor'					):
																				if not isinstance(lowerCAmelCase_			,  lowerCAmelCase_					):
																									module.set_processor(lowerCAmelCase_					)
																				else:
																									module.set_processor(processor.pop(F'''{name}.processor'''					)					)

															for sub_name, child in module.named_children():
																				fn_recursive_attn_processor(F'''{name}.{sub_name}'''			,  lowerCAmelCase_			,  lowerCAmelCase_					)

										for name, module in self.named_children():
															fn_recursive_attn_processor(lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)




					def 		lowerCAmelCase     (						self					)						->							Optional[int]:
										self.set_attn_processor(AttnProcessor()					)




					@apply_forward_hook
					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_ = True					)						->							AutoencoderKLOutput:
										if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
															return self.tiled_encode(lowerCAmelCase_			,  return_dict=lowerCAmelCase_					)

										if self.use_slicing and x.shape[0] > 1:
															_snake_case													=			[self.encoder(lowerCAmelCase_					) for x_slice in x.split(1					)]
															_snake_case													=			torch.cat(lowerCAmelCase_					)
										else:
															_snake_case													=			self.encoder(lowerCAmelCase_					)

										_snake_case													=			self.quant_conv(lowerCAmelCase_					)
										_snake_case													=			DiagonalGaussianDistribution(lowerCAmelCase_					)

										if not return_dict:
															return (posterior,)

										return AutoencoderKLOutput(latent_dist=lowerCAmelCase_					)




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_ = True					)						->							Union[DecoderOutput, torch.FloatTensor]:
										if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
															return self.tiled_decode(lowerCAmelCase_			,  return_dict=lowerCAmelCase_					)

										_snake_case													=			self.post_quant_conv(lowerCAmelCase_					)
										_snake_case													=			self.decoder(lowerCAmelCase_					)

										if not return_dict:
															return (dec,)

										return DecoderOutput(sample=lowerCAmelCase_					)




					@apply_forward_hook
					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_ = True					)						->							Union[DecoderOutput, torch.FloatTensor]:
										if self.use_slicing and z.shape[0] > 1:
															_snake_case													=			[self._decode(lowerCAmelCase_					).sample for z_slice in z.split(1					)]
															_snake_case													=			torch.cat(lowerCAmelCase_					)
										else:
															_snake_case													=			self._decode(lowerCAmelCase_					).sample

										if not return_dict:
															return (decoded,)

										return DecoderOutput(sample=lowerCAmelCase_					)




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)						->							Tuple:
										_snake_case													=			min(a.shape[2]			,  b.shape[2]			,  lowerCAmelCase_					)
										for y in range(lowerCAmelCase_					):
															_snake_case													=			a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
										return b




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)						->							Optional[Any]:
										_snake_case													=			min(a.shape[3]			,  b.shape[3]			,  lowerCAmelCase_					)
										for x in range(lowerCAmelCase_					):
															_snake_case													=			a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
										return b




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_ = True					)						->							AutoencoderKLOutput:
										_snake_case													=			int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)					)
										_snake_case													=			int(self.tile_latent_min_size * self.tile_overlap_factor					)
										_snake_case													=			self.tile_latent_min_size - blend_extent

										# Split the image into 512x512 tiles and encode them separately.
										_snake_case													=			[]
										for i in range(0			,  x.shape[2]			,  lowerCAmelCase_					):
															_snake_case													=			[]
															for j in range(0			,  x.shape[3]			,  lowerCAmelCase_					):
																				_snake_case													=			x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
																				_snake_case													=			self.encoder(lowerCAmelCase_					)
																				_snake_case													=			self.quant_conv(lowerCAmelCase_					)
																				row.append(lowerCAmelCase_					)
															rows.append(lowerCAmelCase_					)
										_snake_case													=			[]
										for i, row in enumerate(lowerCAmelCase_					):
															_snake_case													=			[]
															for j, tile in enumerate(lowerCAmelCase_					):
																				# blend the above tile and the left tile
																				# to the current tile and add the current tile to the result row
																				if i > 0:
																									_snake_case													=			self.blend_v(rows[i - 1][j]			,  lowerCAmelCase_			,  lowerCAmelCase_					)
																				if j > 0:
																									_snake_case													=			self.blend_h(row[j - 1]			,  lowerCAmelCase_			,  lowerCAmelCase_					)
																				result_row.append(tile[:, :, :row_limit, :row_limit]					)
															result_rows.append(torch.cat(lowerCAmelCase_			,  dim=3					)					)

										_snake_case													=			torch.cat(lowerCAmelCase_			,  dim=2					)
										_snake_case													=			DiagonalGaussianDistribution(lowerCAmelCase_					)

										if not return_dict:
															return (posterior,)

										return AutoencoderKLOutput(latent_dist=lowerCAmelCase_					)




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_ = True					)						->							Union[DecoderOutput, torch.FloatTensor]:
										_snake_case													=			int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)					)
										_snake_case													=			int(self.tile_sample_min_size * self.tile_overlap_factor					)
										_snake_case													=			self.tile_sample_min_size - blend_extent

										# Split z into overlapping 64x64 tiles and decode them separately.
										# The tiles have an overlap to avoid seams between tiles.
										_snake_case													=			[]
										for i in range(0			,  z.shape[2]			,  lowerCAmelCase_					):
															_snake_case													=			[]
															for j in range(0			,  z.shape[3]			,  lowerCAmelCase_					):
																				_snake_case													=			z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
																				_snake_case													=			self.post_quant_conv(lowerCAmelCase_					)
																				_snake_case													=			self.decoder(lowerCAmelCase_					)
																				row.append(lowerCAmelCase_					)
															rows.append(lowerCAmelCase_					)
										_snake_case													=			[]
										for i, row in enumerate(lowerCAmelCase_					):
															_snake_case													=			[]
															for j, tile in enumerate(lowerCAmelCase_					):
																				# blend the above tile and the left tile
																				# to the current tile and add the current tile to the result row
																				if i > 0:
																									_snake_case													=			self.blend_v(rows[i - 1][j]			,  lowerCAmelCase_			,  lowerCAmelCase_					)
																				if j > 0:
																									_snake_case													=			self.blend_h(row[j - 1]			,  lowerCAmelCase_			,  lowerCAmelCase_					)
																				result_row.append(tile[:, :, :row_limit, :row_limit]					)
															result_rows.append(torch.cat(lowerCAmelCase_			,  dim=3					)					)

										_snake_case													=			torch.cat(lowerCAmelCase_			,  dim=2					)
										if not return_dict:
															return (dec,)

										return DecoderOutput(sample=lowerCAmelCase_					)






					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_ = False			,  lowerCAmelCase_ = True			,  lowerCAmelCase_ = None			,  )						->							Union[DecoderOutput, torch.FloatTensor]:
										_snake_case													=			sample
										_snake_case													=			self.encode(lowerCAmelCase_					).latent_dist
										if sample_posterior:
															_snake_case													=			posterior.sample(generator=lowerCAmelCase_					)
										else:
															_snake_case													=			posterior.mode()
										_snake_case													=			self.decode(lowerCAmelCase_					).sample

										if not return_dict:
															return (dec,)

										return DecoderOutput(sample=lowerCAmelCase_					)










from typing import Dict, List

from nltk.translate import gleu_score

import datasets
from datasets import MetricInfo


UpperCAmelCase_ 							=       """\
@misc{wu2016googles,
      title={Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation},
      author={Yonghui Wu and Mike Schuster and Zhifeng Chen and Quoc V. Le and Mohammad Norouzi and Wolfgang Macherey
              and Maxim Krikun and Yuan Cao and Qin Gao and Klaus Macherey and Jeff Klingner and Apurva Shah and Melvin
              Johnson and Xiaobing Liu and Łukasz Kaiser and Stephan Gouws and Yoshikiyo Kato and Taku Kudo and Hideto
              Kazawa and Keith Stevens and George Kurian and Nishant Patil and Wei Wang and Cliff Young and
              Jason Smith and Jason Riesa and Alex Rudnick and Oriol Vinyals and Greg Corrado and Macduff Hughes
              and Jeffrey Dean},
      year={2016},
      eprint={1609.08144},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}
"""

UpperCAmelCase_ 							=       """\
The BLEU score has some undesirable properties when used for single
sentences, as it was designed to be a corpus measure. We therefore
use a slightly different score for our RL experiments which we call
the 'GLEU score'. For the GLEU score, we record all sub-sequences of
1, 2, 3 or 4 tokens in output and target sequence (n-grams). We then
compute a recall, which is the ratio of the number of matching n-grams
to the number of total n-grams in the target (ground truth) sequence,
and a precision, which is the ratio of the number of matching n-grams
to the number of total n-grams in the generated output sequence. Then
GLEU score is simply the minimum of recall and precision. This GLEU
score's range is always between 0 (no matches) and 1 (all match) and
it is symmetrical when switching output and target. According to
our experiments, GLEU score correlates quite well with the BLEU
metric on a corpus level but does not have its drawbacks for our per
sentence reward objective.
"""

UpperCAmelCase_ 							=       """\
Computes corpus-level Google BLEU (GLEU) score of translated segments against one or more references.
Instead of averaging the sentence level GLEU scores (i.e. macro-average precision), Wu et al. (2016) sum up the matching
tokens and the max of hypothesis and reference tokens for each sentence, then compute using the aggregate values.

Args:
    predictions (list of str): list of translations to score.
        Each translation should be tokenized into a list of tokens.
    references (list of list of str): list of lists of references for each translation.
        Each reference should be tokenized into a list of tokens.
    min_len (int): The minimum order of n-gram this function should extract. Defaults to 1.
    max_len (int): The maximum order of n-gram this function should extract. Defaults to 4.

Returns:
    'google_bleu': google_bleu score

Examples:
    Example 1:
        >>> hyp1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'which',
        ...         'ensures', 'that', 'the', 'rubber', 'duck', 'always',
        ...         'disobeys', 'the', 'commands', 'of', 'the', 'cat']
        >>> ref1a = ['It', 'is', 'the', 'guiding', 'principle', 'which',
        ...          'guarantees', 'the', 'rubber', 'duck', 'forces', 'never',
        ...          'being', 'under', 'the', 'command', 'of', 'the', 'cat']

        >>> hyp2 = ['he', 'read', 'the', 'book', 'because', 'he', 'was',
        ...         'interested', 'in', 'world', 'history']
        >>> ref2a = ['he', 'was', 'interested', 'in', 'world', 'history',
        ...          'because', 'he', 'read', 'the', 'book']

        >>> list_of_references = [[ref1a], [ref2a]]
        >>> hypotheses = [hyp1, hyp2]
        >>> google_bleu = datasets.load_metric(\"google_bleu\")
        >>> results = google_bleu.compute(predictions=hypotheses, references=list_of_references)
        >>> print(round(results[\"google_bleu\"], 2))
        0.44

    Example 2:
        >>> hyp1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'which',
        ...         'ensures', 'that', 'the', 'rubber', 'duck', 'always',
        ...         'disobeys', 'the', 'commands', 'of', 'the', 'cat']
        >>> ref1a = ['It', 'is', 'the', 'guiding', 'principle', 'which',
        ...          'guarantees', 'the', 'rubber', 'duck', 'forces', 'never',
        ...          'being', 'under', 'the', 'command', 'of', 'the', 'cat']
        >>> ref1b = ['It', 'is', 'a', 'guide', 'to', 'action', 'that',
        ...          'ensures', 'that', 'the', 'rubber', 'duck', 'will', 'never',
        ...          'heed', 'the', 'cat', 'commands']
        >>> ref1c = ['It', 'is', 'the', 'practical', 'guide', 'for', 'the',
        ...          'rubber', 'duck', 'army', 'never', 'to', 'heed', 'the', 'directions',
        ...          'of', 'the', 'cat']

        >>> hyp2 = ['he', 'read', 'the', 'book', 'because', 'he', 'was',
        ...         'interested', 'in', 'world', 'history']
        >>> ref2a = ['he', 'was', 'interested', 'in', 'world', 'history',
        ...          'because', 'he', 'read', 'the', 'book']

        >>> list_of_references = [[ref1a, ref1b, ref1c], [ref2a]]
        >>> hypotheses = [hyp1, hyp2]
        >>> google_bleu = datasets.load_metric(\"google_bleu\")
        >>> results = google_bleu.compute(predictions=hypotheses, references=list_of_references)
        >>> print(round(results[\"google_bleu\"], 2))
        0.61

    Example 3:
        >>> hyp1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'which',
        ...         'ensures', 'that', 'the', 'rubber', 'duck', 'always',
        ...         'disobeys', 'the', 'commands', 'of', 'the', 'cat']
        >>> ref1a = ['It', 'is', 'the', 'guiding', 'principle', 'which',
        ...          'guarantees', 'the', 'rubber', 'duck', 'forces', 'never',
        ...          'being', 'under', 'the', 'command', 'of', 'the', 'cat']
        >>> ref1b = ['It', 'is', 'a', 'guide', 'to', 'action', 'that',
        ...          'ensures', 'that', 'the', 'rubber', 'duck', 'will', 'never',
        ...          'heed', 'the', 'cat', 'commands']
        >>> ref1c = ['It', 'is', 'the', 'practical', 'guide', 'for', 'the',
        ...          'rubber', 'duck', 'army', 'never', 'to', 'heed', 'the', 'directions',
        ...          'of', 'the', 'cat']

        >>> hyp2 = ['he', 'read', 'the', 'book', 'because', 'he', 'was',
        ...         'interested', 'in', 'world', 'history']
        >>> ref2a = ['he', 'was', 'interested', 'in', 'world', 'history',
        ...          'because', 'he', 'read', 'the', 'book']

        >>> list_of_references = [[ref1a, ref1b, ref1c], [ref2a]]
        >>> hypotheses = [hyp1, hyp2]
        >>> google_bleu = datasets.load_metric(\"google_bleu\")
        >>> results = google_bleu.compute(predictions=hypotheses, references=list_of_references, min_len=2)
        >>> print(round(results[\"google_bleu\"], 2))
        0.53

    Example 4:
        >>> hyp1 = ['It', 'is', 'a', 'guide', 'to', 'action', 'which',
        ...         'ensures', 'that', 'the', 'rubber', 'duck', 'always',
        ...         'disobeys', 'the', 'commands', 'of', 'the', 'cat']
        >>> ref1a = ['It', 'is', 'the', 'guiding', 'principle', 'which',
        ...          'guarantees', 'the', 'rubber', 'duck', 'forces', 'never',
        ...          'being', 'under', 'the', 'command', 'of', 'the', 'cat']
        >>> ref1b = ['It', 'is', 'a', 'guide', 'to', 'action', 'that',
        ...          'ensures', 'that', 'the', 'rubber', 'duck', 'will', 'never',
        ...          'heed', 'the', 'cat', 'commands']
        >>> ref1c = ['It', 'is', 'the', 'practical', 'guide', 'for', 'the',
        ...          'rubber', 'duck', 'army', 'never', 'to', 'heed', 'the', 'directions',
        ...          'of', 'the', 'cat']

        >>> hyp2 = ['he', 'read', 'the', 'book', 'because', 'he', 'was',
        ...         'interested', 'in', 'world', 'history']
        >>> ref2a = ['he', 'was', 'interested', 'in', 'world', 'history',
        ...          'because', 'he', 'read', 'the', 'book']

        >>> list_of_references = [[ref1a, ref1b, ref1c], [ref2a]]
        >>> hypotheses = [hyp1, hyp2]
        >>> google_bleu = datasets.load_metric(\"google_bleu\")
        >>> results = google_bleu.compute(predictions=hypotheses,references=list_of_references, min_len=2, max_len=6)
        >>> print(round(results[\"google_bleu\"], 2))
        0.4
"""
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION							,					_KWARGS_DESCRIPTION   )
class      UpperCamelCase_						(					datasets.Metric   ):




					def 		lowerCAmelCase     (						self					)						->							MetricInfo:
										return datasets.MetricInfo(
										    description=_DESCRIPTION			,  citation=_CITATION			,  inputs_description=_KWARGS_DESCRIPTION			,  features=datasets.Features(
										        {
										            'predictions': datasets.Sequence(datasets.Value('string'			,  id='token'					)			,  id='sequence'					),
										            'references': datasets.Sequence(
										                datasets.Sequence(datasets.Value('string'			,  id='token'					)			,  id='sequence'					)			,  id='references'					),
										        }					)			,  )






					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_ = 1			,  lowerCAmelCase_ = 4			,  )						->							Dict[str, float]:
										return {
										    "google_bleu": gleu_score.corpus_gleu(
										        list_of_references=lowerCAmelCase_			,  hypotheses=lowerCAmelCase_			,  min_len=lowerCAmelCase_			,  max_len=lowerCAmelCase_					)
										}



