

"""simple docstring"""





from typing import Dict, List

from nltk.translate import gleu_score

import datasets
from datasets import MetricInfo


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

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

_a							=							'\\nComputes corpus-level Google BLEU (GLEU) score of translated segments against one or more references.\nInstead of averaging the sentence level GLEU scores (i.e. macro-average precision), Wu et al. (2016) sum up the matching\ntokens and the max of hypothesis and reference tokens for each sentence, then compute using the aggregate values.\n\nArgs:\n    predictions (list of str): list of translations to score.\n        Each translation should be tokenized into a list of tokens.\n    references (list of list of str): list of lists of references for each translation.\n        Each reference should be tokenized into a list of tokens.\n    min_len (int): The minimum order of n-gram this function should extract. Defaults to 1.\n    max_len (int): The maximum order of n-gram this function should extract. Defaults to 4.\n\nReturns:\n    \'google_bleu\': google_bleu score\n\nExamples:\n    Example 1:\n        >>> hyp1 = [\'It\', \'is\', \'a\', \'guide\', \'to\', \'action\', \'which\',\n        ...         \'ensures\', \'that\', \'the\', \'rubber\', \'duck\', \'always\',\n        ...         \'disobeys\', \'the\', \'commands\', \'of\', \'the\', \'cat\']\n        >>> ref1a = [\'It\', \'is\', \'the\', \'guiding\', \'principle\', \'which\',\n        ...          \'guarantees\', \'the\', \'rubber\', \'duck\', \'forces\', \'never\',\n        ...          \'being\', \'under\', \'the\', \'command\', \'of\', \'the\', \'cat\']\n\n        >>> hyp2 = [\'he\', \'read\', \'the\', \'book\', \'because\', \'he\', \'was\',\n        ...         \'interested\', \'in\', \'world\', \'history\']\n        >>> ref2a = [\'he\', \'was\', \'interested\', \'in\', \'world\', \'history\',\n        ...          \'because\', \'he\', \'read\', \'the\', \'book\']\n\n        >>> list_of_references = [[ref1a], [ref2a]]\n        >>> hypotheses = [hyp1, hyp2]\n        >>> google_bleu = datasets.load_metric("google_bleu")\n        >>> results = google_bleu.compute(predictions=hypotheses, references=list_of_references)\n        >>> print(round(results["google_bleu"], 2))\n        0.44\n\n    Example 2:\n        >>> hyp1 = [\'It\', \'is\', \'a\', \'guide\', \'to\', \'action\', \'which\',\n        ...         \'ensures\', \'that\', \'the\', \'rubber\', \'duck\', \'always\',\n        ...         \'disobeys\', \'the\', \'commands\', \'of\', \'the\', \'cat\']\n        >>> ref1a = [\'It\', \'is\', \'the\', \'guiding\', \'principle\', \'which\',\n        ...          \'guarantees\', \'the\', \'rubber\', \'duck\', \'forces\', \'never\',\n        ...          \'being\', \'under\', \'the\', \'command\', \'of\', \'the\', \'cat\']\n        >>> ref1b = [\'It\', \'is\', \'a\', \'guide\', \'to\', \'action\', \'that\',\n        ...          \'ensures\', \'that\', \'the\', \'rubber\', \'duck\', \'will\', \'never\',\n        ...          \'heed\', \'the\', \'cat\', \'commands\']\n        >>> ref1c = [\'It\', \'is\', \'the\', \'practical\', \'guide\', \'for\', \'the\',\n        ...          \'rubber\', \'duck\', \'army\', \'never\', \'to\', \'heed\', \'the\', \'directions\',\n        ...          \'of\', \'the\', \'cat\']\n\n        >>> hyp2 = [\'he\', \'read\', \'the\', \'book\', \'because\', \'he\', \'was\',\n        ...         \'interested\', \'in\', \'world\', \'history\']\n        >>> ref2a = [\'he\', \'was\', \'interested\', \'in\', \'world\', \'history\',\n        ...          \'because\', \'he\', \'read\', \'the\', \'book\']\n\n        >>> list_of_references = [[ref1a, ref1b, ref1c], [ref2a]]\n        >>> hypotheses = [hyp1, hyp2]\n        >>> google_bleu = datasets.load_metric("google_bleu")\n        >>> results = google_bleu.compute(predictions=hypotheses, references=list_of_references)\n        >>> print(round(results["google_bleu"], 2))\n        0.61\n\n    Example 3:\n        >>> hyp1 = [\'It\', \'is\', \'a\', \'guide\', \'to\', \'action\', \'which\',\n        ...         \'ensures\', \'that\', \'the\', \'rubber\', \'duck\', \'always\',\n        ...         \'disobeys\', \'the\', \'commands\', \'of\', \'the\', \'cat\']\n        >>> ref1a = [\'It\', \'is\', \'the\', \'guiding\', \'principle\', \'which\',\n        ...          \'guarantees\', \'the\', \'rubber\', \'duck\', \'forces\', \'never\',\n        ...          \'being\', \'under\', \'the\', \'command\', \'of\', \'the\', \'cat\']\n        >>> ref1b = [\'It\', \'is\', \'a\', \'guide\', \'to\', \'action\', \'that\',\n        ...          \'ensures\', \'that\', \'the\', \'rubber\', \'duck\', \'will\', \'never\',\n        ...          \'heed\', \'the\', \'cat\', \'commands\']\n        >>> ref1c = [\'It\', \'is\', \'the\', \'practical\', \'guide\', \'for\', \'the\',\n        ...          \'rubber\', \'duck\', \'army\', \'never\', \'to\', \'heed\', \'the\', \'directions\',\n        ...          \'of\', \'the\', \'cat\']\n\n        >>> hyp2 = [\'he\', \'read\', \'the\', \'book\', \'because\', \'he\', \'was\',\n        ...         \'interested\', \'in\', \'world\', \'history\']\n        >>> ref2a = [\'he\', \'was\', \'interested\', \'in\', \'world\', \'history\',\n        ...          \'because\', \'he\', \'read\', \'the\', \'book\']\n\n        >>> list_of_references = [[ref1a, ref1b, ref1c], [ref2a]]\n        >>> hypotheses = [hyp1, hyp2]\n        >>> google_bleu = datasets.load_metric("google_bleu")\n        >>> results = google_bleu.compute(predictions=hypotheses, references=list_of_references, min_len=2)\n        >>> print(round(results["google_bleu"], 2))\n        0.53\n\n    Example 4:\n        >>> hyp1 = [\'It\', \'is\', \'a\', \'guide\', \'to\', \'action\', \'which\',\n        ...         \'ensures\', \'that\', \'the\', \'rubber\', \'duck\', \'always\',\n        ...         \'disobeys\', \'the\', \'commands\', \'of\', \'the\', \'cat\']\n        >>> ref1a = [\'It\', \'is\', \'the\', \'guiding\', \'principle\', \'which\',\n        ...          \'guarantees\', \'the\', \'rubber\', \'duck\', \'forces\', \'never\',\n        ...          \'being\', \'under\', \'the\', \'command\', \'of\', \'the\', \'cat\']\n        >>> ref1b = [\'It\', \'is\', \'a\', \'guide\', \'to\', \'action\', \'that\',\n        ...          \'ensures\', \'that\', \'the\', \'rubber\', \'duck\', \'will\', \'never\',\n        ...          \'heed\', \'the\', \'cat\', \'commands\']\n        >>> ref1c = [\'It\', \'is\', \'the\', \'practical\', \'guide\', \'for\', \'the\',\n        ...          \'rubber\', \'duck\', \'army\', \'never\', \'to\', \'heed\', \'the\', \'directions\',\n        ...          \'of\', \'the\', \'cat\']\n\n        >>> hyp2 = [\'he\', \'read\', \'the\', \'book\', \'because\', \'he\', \'was\',\n        ...         \'interested\', \'in\', \'world\', \'history\']\n        >>> ref2a = [\'he\', \'was\', \'interested\', \'in\', \'world\', \'history\',\n        ...          \'because\', \'he\', \'read\', \'the\', \'book\']\n\n        >>> list_of_references = [[ref1a, ref1b, ref1c], [ref2a]]\n        >>> hypotheses = [hyp1, hyp2]\n        >>> google_bleu = datasets.load_metric("google_bleu")\n        >>> results = google_bleu.compute(predictions=hypotheses,references=list_of_references, min_len=2, max_len=6)\n        >>> print(round(results["google_bleu"], 2))\n        0.4\n'




@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION					,_KWARGS_DESCRIPTION       )
class  _lowerCAmelCase  (     datasets.Metric       ):







 """simple docstring"""







 def  _lowercase			(      self							:      Union[str, Any]     ):
    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  _lowercase			(      self							:      Dict, UpperCAmelCase__							:      List[List[List[str]]], UpperCAmelCase__							:      List[List[str]], UpperCAmelCase__							:      int = 1, UpperCAmelCase__							:      int = 4, ):
    return {
        "google_bleu": gleu_score.corpus_gleu(
            list_of_references=UpperCAmelCase__, hypotheses=UpperCAmelCase__, min_len=UpperCAmelCase__, max_len=UpperCAmelCase__     )
    }







"""simple docstring"""





from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)




@flax.struct.dataclass
class  _lowerCAmelCase  (     lowercase       ):







 """simple docstring"""







 __UpperCAmelCase     : jnp.ndarray
 __UpperCAmelCase     : jnp.ndarray







class  _lowerCAmelCase  (     nn.Module       ):







 """simple docstring"""







 __UpperCAmelCase     : int
 __UpperCAmelCase     : Tuple[int]       		=     (1_6, 3_2, 9_6, 2_5_6)
 __UpperCAmelCase     : jnp.dtype       		=     jnp.floataa


 def  _lowercase			(      self							:      Union[str, Any]     ):
    __lowercase								=						nn.Conv(
        self.block_out_channels[0], kernel_size=(3, 3), padding=((1, 1), (1, 1)), dtype=self.dtype, )

    __lowercase								=						[]
    for i in range(len(self.block_out_channels     ) - 1     ):
       __lowercase								=						self.block_out_channels[i]
       __lowercase								=						self.block_out_channels[i + 1]
       __lowercase								=						nn.Conv(
           UpperCAmelCase__, kernel_size=(3, 3), padding=((1, 1), (1, 1)), dtype=self.dtype, )
       blocks.append(UpperCAmelCase__     )
       __lowercase								=						nn.Conv(
           UpperCAmelCase__, kernel_size=(3, 3), strides=(2, 2), padding=((1, 1), (1, 1)), dtype=self.dtype, )
       blocks.append(UpperCAmelCase__     )
    __lowercase								=						blocks

    __lowercase								=						nn.Conv(
        self.conditioning_embedding_channels, kernel_size=(3, 3), padding=((1, 1), (1, 1)), kernel_init=nn.initializers.zeros_init(), bias_init=nn.initializers.zeros_init(), dtype=self.dtype, )



 def __call__(      self							:      Any, UpperCAmelCase__							:      Union[str, Any]     ):
    __lowercase								=						self.conv_in(UpperCAmelCase__     )
    __lowercase								=						nn.silu(UpperCAmelCase__     )

    for block in self.blocks:
       __lowercase								=						block(UpperCAmelCase__     )
       __lowercase								=						nn.silu(UpperCAmelCase__     )

    __lowercase								=						self.conv_out(UpperCAmelCase__     )

    return embedding







@flax_register_to_config
class  _lowerCAmelCase  (     nn.Module					,lowercase					,lowercase       ):







 """simple docstring"""







 __UpperCAmelCase     : int       		=     3_2
 __UpperCAmelCase     : int       		=     4
 __UpperCAmelCase     : Tuple[str]       		=     (
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "DownBlock2D",
 )
 __UpperCAmelCase     : Union[bool, Tuple[bool]]       		=     False
 __UpperCAmelCase     : Tuple[int]       		=     (3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
 __UpperCAmelCase     : int       		=     2
 __UpperCAmelCase     : Union[int, Tuple[int]]       		=     8
 __UpperCAmelCase     : Optional[Union[int, Tuple[int]]]       		=     None
 __UpperCAmelCase     : int       		=     1_2_8_0
 __UpperCAmelCase     : float       		=     0.0
 __UpperCAmelCase     : bool       		=     False
 __UpperCAmelCase     : jnp.dtype       		=     jnp.floataa
 __UpperCAmelCase     : bool       		=     True
 __UpperCAmelCase     : int       		=     0
 __UpperCAmelCase     : str       		=     "rgb"
 __UpperCAmelCase     : Tuple[int]       		=     (1_6, 3_2, 9_6, 2_5_6)


 def  _lowercase			(      self							:      List[Any], UpperCAmelCase__							:      jax.random.KeyArray     ):
    # init input tensors
    __lowercase								=						(1, self.in_channels, self.sample_size, self.sample_size)
    __lowercase								=						jnp.zeros(UpperCAmelCase__, dtype=jnp.floataa     )
    __lowercase								=						jnp.ones((1,), dtype=jnp.intaa     )
    __lowercase								=						jnp.zeros((1, 1, self.cross_attention_dim), dtype=jnp.floataa     )
    __lowercase								=						(1, 3, self.sample_size * 8, self.sample_size * 8)
    __lowercase								=						jnp.zeros(UpperCAmelCase__, dtype=jnp.floataa     )

    __lowercase		,__lowercase								=						jax.random.split(UpperCAmelCase__     )
    __lowercase								=						{"params": params_rng, "dropout": dropout_rng}

    return self.init(UpperCAmelCase__, UpperCAmelCase__, UpperCAmelCase__, UpperCAmelCase__, UpperCAmelCase__     )["params"]


 def  _lowercase			(      self							:      Union[str, Any]     ):
    __lowercase								=						self.block_out_channels
    __lowercase								=						block_out_channels[0] * 4

    # If `num_attention_heads` is not defined (which is the case for most models)
    # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
    # The reason for this behavior is to correct for incorrectly named variables that were introduced
    # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
    # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
    # which is why we correct for the naming here.
    __lowercase								=						self.num_attention_heads or self.attention_head_dim

    # input
    __lowercase								=						nn.Conv(
        block_out_channels[0], kernel_size=(3, 3), strides=(1, 1), padding=((1, 1), (1, 1)), dtype=self.dtype, )

    # time
    __lowercase								=						FlaxTimesteps(
        block_out_channels[0], flip_sin_to_cos=self.flip_sin_to_cos, freq_shift=self.config.freq_shift     )
    __lowercase								=						FlaxTimestepEmbedding(UpperCAmelCase__, dtype=self.dtype     )

    __lowercase								=						FlaxControlNetConditioningEmbedding(
        conditioning_embedding_channels=block_out_channels[0], block_out_channels=self.conditioning_embedding_out_channels, )

    __lowercase								=						self.only_cross_attention
    if isinstance(UpperCAmelCase__, UpperCAmelCase__     ):
       __lowercase								=						(only_cross_attention,) * len(self.down_block_types     )

    if isinstance(UpperCAmelCase__, UpperCAmelCase__     ):
       __lowercase								=						(num_attention_heads,) * len(self.down_block_types     )

    # down
    __lowercase								=						[]
    __lowercase								=						[]

    __lowercase								=						block_out_channels[0]

    __lowercase								=						nn.Conv(
        UpperCAmelCase__, kernel_size=(1, 1), padding="VALID", kernel_init=nn.initializers.zeros_init(), bias_init=nn.initializers.zeros_init(), dtype=self.dtype, )
    controlnet_down_blocks.append(UpperCAmelCase__     )

    for i, down_block_type in enumerate(self.down_block_types     ):
       __lowercase								=						output_channel
       __lowercase								=						block_out_channels[i]
       __lowercase								=						i == len(UpperCAmelCase__     ) - 1

       if down_block_type == "CrossAttnDownBlock2D":
          __lowercase								=						FlaxCrossAttnDownBlockaD(
              in_channels=UpperCAmelCase__, out_channels=UpperCAmelCase__, dropout=self.dropout, num_layers=self.layers_per_block, num_attention_heads=num_attention_heads[i], add_downsample=not is_final_block, use_linear_projection=self.use_linear_projection, only_cross_attention=only_cross_attention[i], dtype=self.dtype, )
       else:
          __lowercase								=						FlaxDownBlockaD(
              in_channels=UpperCAmelCase__, out_channels=UpperCAmelCase__, dropout=self.dropout, num_layers=self.layers_per_block, add_downsample=not is_final_block, dtype=self.dtype, )

       down_blocks.append(UpperCAmelCase__     )

       for _ in range(self.layers_per_block     ):
          __lowercase								=						nn.Conv(
              UpperCAmelCase__, kernel_size=(1, 1), padding="VALID", kernel_init=nn.initializers.zeros_init(), bias_init=nn.initializers.zeros_init(), dtype=self.dtype, )
          controlnet_down_blocks.append(UpperCAmelCase__     )

       if not is_final_block:
          __lowercase								=						nn.Conv(
              UpperCAmelCase__, kernel_size=(1, 1), padding="VALID", kernel_init=nn.initializers.zeros_init(), bias_init=nn.initializers.zeros_init(), dtype=self.dtype, )
          controlnet_down_blocks.append(UpperCAmelCase__     )

    __lowercase								=						down_blocks
    __lowercase								=						controlnet_down_blocks

    # mid
    __lowercase								=						block_out_channels[-1]
    __lowercase								=						FlaxUNetMidBlockaDCrossAttn(
        in_channels=UpperCAmelCase__, dropout=self.dropout, num_attention_heads=num_attention_heads[-1], use_linear_projection=self.use_linear_projection, dtype=self.dtype, )

    __lowercase								=						nn.Conv(
        UpperCAmelCase__, kernel_size=(1, 1), padding="VALID", kernel_init=nn.initializers.zeros_init(), bias_init=nn.initializers.zeros_init(), dtype=self.dtype, )



 def __call__(      self							:      str, UpperCAmelCase__							:      Union[str, Any], UpperCAmelCase__							:      List[Any], UpperCAmelCase__							:      Tuple, UpperCAmelCase__							:      List[str], UpperCAmelCase__							:      float = 1.0, UpperCAmelCase__							:      bool = True, UpperCAmelCase__							:      bool = False, ):
    __lowercase								=						self.controlnet_conditioning_channel_order
    if channel_order == "bgr":
       __lowercase								=						jnp.flip(UpperCAmelCase__, axis=1     )

    # 1. time
    if not isinstance(UpperCAmelCase__, jnp.ndarray     ):
       __lowercase								=						jnp.array([timesteps], dtype=jnp.intaa     )
    elif isinstance(UpperCAmelCase__, jnp.ndarray     ) and len(timesteps.shape     ) == 0:
       __lowercase								=						timesteps.astype(dtype=jnp.floataa     )
       __lowercase								=						jnp.expand_dims(UpperCAmelCase__, 0     )

    __lowercase								=						self.time_proj(UpperCAmelCase__     )
    __lowercase								=						self.time_embedding(UpperCAmelCase__     )

    # 2. pre-process
    __lowercase								=						jnp.transpose(UpperCAmelCase__, (0, 2, 3, 1)     )
    __lowercase								=						self.conv_in(UpperCAmelCase__     )

    __lowercase								=						jnp.transpose(UpperCAmelCase__, (0, 2, 3, 1)     )
    __lowercase								=						self.controlnet_cond_embedding(UpperCAmelCase__     )
    sample += controlnet_cond

    # 3. down
    __lowercase								=						(sample,)
    for down_block in self.down_blocks:
       if isinstance(UpperCAmelCase__, UpperCAmelCase__     ):
          __lowercase		,__lowercase								=						down_block(UpperCAmelCase__, UpperCAmelCase__, UpperCAmelCase__, deterministic=not train     )
       else:
          __lowercase		,__lowercase								=						down_block(UpperCAmelCase__, UpperCAmelCase__, deterministic=not train     )
       down_block_res_samples += res_samples

    # 4. mid
    __lowercase								=						self.mid_block(UpperCAmelCase__, UpperCAmelCase__, UpperCAmelCase__, deterministic=not train     )

    # 5. contronet blocks
    __lowercase								=						()
    for down_block_res_sample, controlnet_block in zip(UpperCAmelCase__, self.controlnet_down_blocks     ):
       __lowercase								=						controlnet_block(UpperCAmelCase__     )
       controlnet_down_block_res_samples += (down_block_res_sample,)

    __lowercase								=						controlnet_down_block_res_samples

    __lowercase								=						self.controlnet_mid_block(UpperCAmelCase__     )

    # 6. scaling
    __lowercase								=						[sample * conditioning_scale for sample in down_block_res_samples]
    mid_block_res_sample *= conditioning_scale

    if not return_dict:
       return (down_block_res_samples, mid_block_res_sample)

    return FlaxControlNetOutput(
        down_block_res_samples=UpperCAmelCase__, mid_block_res_sample=UpperCAmelCase__     )




