





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,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)






@flax.struct.dataclass
class        snake_case_    (		__lowercase						):
  A_      							=   42






@flax_register_to_config
class        snake_case_    (		nn.Module					,__lowercase					,__lowercase						):
  A_      							=   32
  A_      							=   4
  A_      							=   4
  A_      							=   (
      "CrossAttnDownBlock2D",
      "CrossAttnDownBlock2D",
      "CrossAttnDownBlock2D",
      "DownBlock2D",
  )
  A_      							=   ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
  A_      							=   False
  A_      							=   (320, 640, 1280, 1280)
  A_      							=   2
  A_      							=   8
  A_      							=   None
  A_      							=   1280
  A_      							=   0.0
  A_      							=   False
  A_      							=   jnp.floataa
  A_      							=   True
  A_      							=   0
  A_      							=   False






  def 				UpperCAmelCase__						(							self			:					Dict       ,  _snake_case			:					jax.random.KeyArray      )->FrozenDict:






   '''simple docstring'''

   __lowerCAmelCase					:	Tuple      =       (1, self.in_channels, self.sample_size, self.sample_size)
   __lowerCAmelCase					:	Union[str, Any]      =       jnp.zeros(_snake_case       ,  dtype=jnp.floataa      )
   __lowerCAmelCase					:	str      =       jnp.ones((1,)       ,  dtype=jnp.intaa      )
   __lowerCAmelCase					:	Dict      =       jnp.zeros((1, 1, self.cross_attention_dim)       ,  dtype=jnp.floataa      )

   __lowerCAmelCase ,					__lowerCAmelCase					:	List[str]      =       jax.random.split(_snake_case      )
   __lowerCAmelCase					:	int      =       {"""params""": params_rng, """dropout""": dropout_rng}

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






  def 				UpperCAmelCase__						(							self			:					Optional[Any]      )->List[str]:






   '''simple docstring'''

   __lowerCAmelCase					:	Union[str, Any]      =       self.block_out_channels
   __lowerCAmelCase					:	Union[str, Any]      =       block_out_channels[0] * 4

   if self.num_attention_heads is not None:
    raise ValueError(
        """At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19."""      )

   # 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.
   __lowerCAmelCase					:	Tuple      =       self.num_attention_heads or self.attention_head_dim

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

   # time
   __lowerCAmelCase					:	List[str]      =       FlaxTimesteps(
       block_out_channels[0]       ,  flip_sin_to_cos=self.flip_sin_to_cos       ,  freq_shift=self.config.freq_shift      )
   __lowerCAmelCase					:	List[Any]      =       FlaxTimestepEmbedding(_snake_case       ,  dtype=self.dtype      )

   __lowerCAmelCase					:	Dict      =       self.only_cross_attention
   if isinstance(_snake_case       ,  _snake_case      ):
    __lowerCAmelCase					:	Union[str, Any]      =       (only_cross_attention,) * len(self.down_block_types      )

   if isinstance(_snake_case       ,  _snake_case      ):
    __lowerCAmelCase					:	str      =       (num_attention_heads,) * len(self.down_block_types      )

   # down
   __lowerCAmelCase					:	Union[str, Any]      =       []
   __lowerCAmelCase					:	Union[str, Any]      =       block_out_channels[0]
   for i, down_block_type in enumerate(self.down_block_types      ):
    __lowerCAmelCase					:	Optional[Any]      =       output_channel
    __lowerCAmelCase					:	Optional[Any]      =       block_out_channels[i]
    __lowerCAmelCase					:	str      =       i == len(_snake_case      ) - 1

    if down_block_type == "CrossAttnDownBlock2D":
     __lowerCAmelCase					:	int      =       FlaxCrossAttnDownBlockaD(
         in_channels=_snake_case       ,  out_channels=_snake_case       ,  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]       ,  use_memory_efficient_attention=self.use_memory_efficient_attention       ,  dtype=self.dtype       ,  )
    else:
     __lowerCAmelCase					:	int      =       FlaxDownBlockaD(
         in_channels=_snake_case       ,  out_channels=_snake_case       ,  dropout=self.dropout       ,  num_layers=self.layers_per_block       ,  add_downsample=not is_final_block       ,  dtype=self.dtype       ,  )

    down_blocks.append(_snake_case      )
   __lowerCAmelCase					:	Union[str, Any]      =       down_blocks

   # mid
   __lowerCAmelCase					:	List[Any]      =       FlaxUNetMidBlockaDCrossAttn(
       in_channels=block_out_channels[-1]       ,  dropout=self.dropout       ,  num_attention_heads=num_attention_heads[-1]       ,  use_linear_projection=self.use_linear_projection       ,  use_memory_efficient_attention=self.use_memory_efficient_attention       ,  dtype=self.dtype       ,  )

   # up
   __lowerCAmelCase					:	Union[str, Any]      =       []
   __lowerCAmelCase					:	Union[str, Any]      =       list(reversed(_snake_case      )      )
   __lowerCAmelCase					:	str      =       list(reversed(_snake_case      )      )
   __lowerCAmelCase					:	Union[str, Any]      =       list(reversed(_snake_case      )      )
   __lowerCAmelCase					:	Any      =       reversed_block_out_channels[0]
   for i, up_block_type in enumerate(self.up_block_types      ):
    __lowerCAmelCase					:	Dict      =       output_channel
    __lowerCAmelCase					:	str      =       reversed_block_out_channels[i]
    __lowerCAmelCase					:	int      =       reversed_block_out_channels[min(i + 1       ,  len(_snake_case      ) - 1      )]

    __lowerCAmelCase					:	str      =       i == len(_snake_case      ) - 1

    if up_block_type == "CrossAttnUpBlock2D":
     __lowerCAmelCase					:	str      =       FlaxCrossAttnUpBlockaD(
         in_channels=_snake_case       ,  out_channels=_snake_case       ,  prev_output_channel=_snake_case       ,  num_layers=self.layers_per_block + 1       ,  num_attention_heads=reversed_num_attention_heads[i]       ,  add_upsample=not is_final_block       ,  dropout=self.dropout       ,  use_linear_projection=self.use_linear_projection       ,  only_cross_attention=only_cross_attention[i]       ,  use_memory_efficient_attention=self.use_memory_efficient_attention       ,  dtype=self.dtype       ,  )
    else:
     __lowerCAmelCase					:	Tuple      =       FlaxUpBlockaD(
         in_channels=_snake_case       ,  out_channels=_snake_case       ,  prev_output_channel=_snake_case       ,  num_layers=self.layers_per_block + 1       ,  add_upsample=not is_final_block       ,  dropout=self.dropout       ,  dtype=self.dtype       ,  )

    up_blocks.append(_snake_case      )
    __lowerCAmelCase					:	Tuple      =       output_channel
   __lowerCAmelCase					:	int      =       up_blocks

   # out
   __lowerCAmelCase					:	Optional[int]      =       nn.GroupNorm(num_groups=32       ,  epsilon=1E-5      )
   __lowerCAmelCase					:	Tuple      =       nn.Conv(
       self.out_channels       ,  kernel_size=(3, 3)       ,  strides=(1, 1)       ,  padding=((1, 1), (1, 1))       ,  dtype=self.dtype       ,  )






  def __call__(							self			:					Any       ,  _snake_case			:					List[Any]       ,  _snake_case			:					int       ,  _snake_case			:					List[Any]       ,  _snake_case			:					str=None       ,  _snake_case			:					str=None       ,  _snake_case			:					bool = True       ,  _snake_case			:					bool = False       ,  )->Union[FlaxUNetaDConditionOutput, Tuple]:






   '''simple docstring'''

   if not isinstance(_snake_case       ,  jnp.ndarray      ):
    __lowerCAmelCase					:	Union[str, Any]      =       jnp.array([timesteps]       ,  dtype=jnp.intaa      )
   elif isinstance(_snake_case       ,  jnp.ndarray      ) and len(timesteps.shape      ) == 0:
    __lowerCAmelCase					:	Union[str, Any]      =       timesteps.astype(dtype=jnp.floataa      )
    __lowerCAmelCase					:	List[Any]      =       jnp.expand_dims(_snake_case       ,  0      )

   __lowerCAmelCase					:	Union[str, Any]      =       self.time_proj(_snake_case      )
   __lowerCAmelCase					:	Tuple      =       self.time_embedding(_snake_case      )

   # 2. pre-process
   __lowerCAmelCase					:	Dict      =       jnp.transpose(_snake_case       ,  (0, 2, 3, 1)      )
   __lowerCAmelCase					:	Tuple      =       self.conv_in(_snake_case      )

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

   if down_block_additional_residuals is not None:
    __lowerCAmelCase					:	Union[str, Any]      =       ()

    for down_block_res_sample, down_block_additional_residual in zip(
        _snake_case       ,  _snake_case      ):
     down_block_res_sample += down_block_additional_residual
     new_down_block_res_samples += (down_block_res_sample,)

    __lowerCAmelCase					:	Dict      =       new_down_block_res_samples

   # 4. mid
   __lowerCAmelCase					:	Any      =       self.mid_block(_snake_case       ,  _snake_case       ,  _snake_case       ,  deterministic=not train      )

   if mid_block_additional_residual is not None:
    sample += mid_block_additional_residual

   # 5. up
   for up_block in self.up_blocks:
    __lowerCAmelCase					:	Tuple      =       down_block_res_samples[-(self.layers_per_block + 1) :]
    __lowerCAmelCase					:	Any      =       down_block_res_samples[: -(self.layers_per_block + 1)]
    if isinstance(_snake_case       ,  _snake_case      ):
     __lowerCAmelCase					:	Optional[int]      =       up_block(
         _snake_case       ,  temb=_snake_case       ,  encoder_hidden_states=_snake_case       ,  res_hidden_states_tuple=_snake_case       ,  deterministic=not train       ,  )
    else:
     __lowerCAmelCase					:	List[Any]      =       up_block(_snake_case       ,  temb=_snake_case       ,  res_hidden_states_tuple=_snake_case       ,  deterministic=not train      )

        # 6. post-process
   __lowerCAmelCase					:	Union[str, Any]      =       self.conv_norm_out(_snake_case      )
   __lowerCAmelCase					:	Optional[Any]      =       nn.silu(_snake_case      )
   __lowerCAmelCase					:	str      =       self.conv_out(_snake_case      )
   __lowerCAmelCase					:	Optional[int]      =       jnp.transpose(_snake_case       ,  (0, 3, 1, 2)      )

   if not return_dict:
    return (sample,)

   return FlaxUNetaDConditionOutput(sample=_snake_case      )






import sacrebleu as scb
from packaging import version
from sacrebleu import TER

import datasets


_UpperCAmelCase             =							'\\n@inproceedings{snover-etal-2006-study,\n    title = "A Study of Translation Edit Rate with Targeted Human Annotation",\n    author = "Snover, Matthew  and\n      Dorr, Bonnie  and\n      Schwartz, Rich  and\n      Micciulla, Linnea  and\n      Makhoul, John",\n    booktitle = "Proceedings of the 7th Conference of the Association for Machine Translation in the Americas: Technical Papers",\n    month = aug # " 8-12",\n    year = "2006",\n    address = "Cambridge, Massachusetts, USA",\n    publisher = "Association for Machine Translation in the Americas",\n    url = "https://aclanthology.org/2006.amta-papers.25",\n    pages = "223--231",\n}\n@inproceedings{post-2018-call,\n    title = "A Call for Clarity in Reporting {BLEU} Scores",\n    author = "Post, Matt",\n    booktitle = "Proceedings of the Third Conference on Machine Translation: Research Papers",\n    month = oct,\n    year = "2018",\n    address = "Belgium, Brussels",\n    publisher = "Association for Computational Linguistics",\n    url = "https://www.aclweb.org/anthology/W18-6319",\n    pages = "186--191",\n}\n'

_UpperCAmelCase             =							'\\nTER (Translation Edit Rate, also called Translation Error Rate) is a metric to quantify the edit operations that a\nhypothesis requires to match a reference translation. We use the implementation that is already present in sacrebleu\n(https://github.com/mjpost/sacreBLEU#ter), which in turn is inspired by the TERCOM implementation, which can be found\nhere: https://github.com/jhclark/tercom.\n\nThe implementation here is slightly different from sacrebleu in terms of the required input format. The length of\nthe references and hypotheses lists need to be the same, so you may need to transpose your references compared to\nsacrebleu\'s required input format. See https://github.com/huggingface/datasets/issues/3154#issuecomment-950746534\n\nSee the README.md file at https://github.com/mjpost/sacreBLEU#ter for more information.\n'

_UpperCAmelCase             =							'\nProduces TER scores alongside the number of edits and reference length.\n\nArgs:\n    predictions (list of str): The system stream (a sequence of segments).\n    references (list of list of str): A list of one or more reference streams (each a sequence of segments).\n    normalized (boolean): If `True`, applies basic tokenization and normalization to sentences. Defaults to `False`.\n    ignore_punct (boolean): If `True`, applies basic tokenization and normalization to sentences. Defaults to `False`.\n    support_zh_ja_chars (boolean): If `True`, tokenization/normalization supports processing of Chinese characters,\n                                    as well as Japanese Kanji, Hiragana, Katakana, and Phonetic Extensions of Katakana.\n                                    Only applies if `normalized = True`. Defaults to `False`.\n    case_sensitive (boolean): If `False`, makes all predictions and references lowercase to ignore differences in case. Defaults to `False`.\n\nReturns:\n    \'score\' (float): TER score (num_edits / sum_ref_lengths * 100)\n    \'num_edits\' (int): The cumulative number of edits\n    \'ref_length\' (float): The cumulative average reference length\n\nExamples:\n    Example 1:\n        >>> predictions = ["does this sentence match??",\n        ...                     "what about this sentence?",\n        ...                     "What did the TER metric user say to the developer?"]\n        >>> references = [["does this sentence match", "does this sentence match!?!"],\n        ...             ["wHaT aBoUt ThIs SeNtEnCe?", "wHaT aBoUt ThIs SeNtEnCe?"],\n        ...             ["Your jokes are...", "...TERrible"]]\n        >>> ter = datasets.load_metric("ter")\n        >>> results = ter.compute(predictions=predictions,\n        ...                         references=references,\n        ...                         case_sensitive=True)\n        >>> print(results)\n        {\'score\': 150.0, \'num_edits\': 15, \'ref_length\': 10.0}\n\n    Example 2:\n        >>> predictions = ["does this sentence match??",\n        ...                     "what about this sentence?"]\n        >>> references = [["does this sentence match", "does this sentence match!?!"],\n        ...             ["wHaT aBoUt ThIs SeNtEnCe?", "wHaT aBoUt ThIs SeNtEnCe?"]]\n        >>> ter = datasets.load_metric("ter")\n        >>> results = ter.compute(predictions=predictions,\n        ...                         references=references,\n        ...                         case_sensitive=True)\n        >>> print(results)\n        {\'score\': 62.5, \'num_edits\': 5, \'ref_length\': 8.0}\n\n    Example 3:\n        >>> predictions = ["does this sentence match??",\n        ...                     "what about this sentence?"]\n        >>> references = [["does this sentence match", "does this sentence match!?!"],\n        ...             ["wHaT aBoUt ThIs SeNtEnCe?", "wHaT aBoUt ThIs SeNtEnCe?"]]\n        >>> ter = datasets.load_metric("ter")\n        >>> results = ter.compute(predictions=predictions,\n        ...                         references=references,\n        ...                         normalized=True,\n        ...                         case_sensitive=True)\n        >>> print(results)\n        {\'score\': 57.14285714285714, \'num_edits\': 6, \'ref_length\': 10.5}\n\n    Example 4:\n        >>> predictions = ["does this sentence match??",\n        ...                     "what about this sentence?"]\n        >>> references = [["does this sentence match", "does this sentence match!?!"],\n        ...             ["wHaT aBoUt ThIs SeNtEnCe?", "wHaT aBoUt ThIs SeNtEnCe?"]]\n        >>> ter = datasets.load_metric("ter")\n        >>> results = ter.compute(predictions=predictions,\n        ...                         references=references,\n        ...                         ignore_punct=True,\n        ...                         case_sensitive=False)\n        >>> print(results)\n        {\'score\': 0.0, \'num_edits\': 0, \'ref_length\': 8.0}\n\n    Example 5:\n        >>> predictions = ["does this sentence match??",\n        ...                    "what about this sentence?",\n        ...                    "What did the TER metric user say to the developer?"]\n        >>> references = [["does this sentence match", "does this sentence match!?!"],\n        ...             ["wHaT aBoUt ThIs SeNtEnCe?", "wHaT aBoUt ThIs SeNtEnCe?"],\n        ...             ["Your jokes are...", "...TERrible"]]\n        >>> ter = datasets.load_metric("ter")\n        >>> results = ter.compute(predictions=predictions,\n        ...                         references=references,\n        ...                         ignore_punct=True,\n        ...                         case_sensitive=False)\n        >>> print(results)\n        {\'score\': 100.0, \'num_edits\': 10, \'ref_length\': 10.0}\n'






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






  def 				UpperCAmelCase__						(							self			:					Tuple      )->str:






   '''simple docstring'''

   if version.parse(scb.__version__      ) < version.parse("""1.4.12"""      ):
    raise ImportWarning(
        """To use `sacrebleu`, the module `sacrebleu>=1.4.12` is required, and the current version of `sacrebleu` doesn't match this condition.\n"""
        """You can install it with `pip install \"sacrebleu>=1.4.12\"`."""      )
   return datasets.MetricInfo(
       description=_DESCRIPTION       ,  citation=_CITATION       ,  homepage="""http://www.cs.umd.edu/~snover/tercom/"""       ,  inputs_description=_KWARGS_DESCRIPTION       ,  features=datasets.Features(
           {
               """predictions""": datasets.Value("""string"""       ,  id="""sequence"""      ),
               """references""": datasets.Sequence(datasets.Value("""string"""       ,  id="""sequence"""      )       ,  id="""references"""      ),
           }      )       ,  codebase_urls=["""https://github.com/mjpost/sacreBLEU#ter"""]       ,  reference_urls=[
           """https://github.com/jhclark/tercom""",
       ]       ,  )






  def 				UpperCAmelCase__						(							self			:					Dict       ,  _snake_case			:					List[Any]       ,  _snake_case			:					int       ,  _snake_case			:					bool = False       ,  _snake_case			:					bool = False       ,  _snake_case			:					bool = False       ,  _snake_case			:					bool = False       ,  )->str:






   '''simple docstring'''

   __lowerCAmelCase					:	List[str]      =       len(references[0]      )
   if any(len(_snake_case      ) != references_per_prediction for refs in references      ):
    raise ValueError("""Sacrebleu requires the same number of references for each prediction"""      )
   __lowerCAmelCase					:	str      =       [[refs[i] for refs in references] for i in range(_snake_case      )]

   __lowerCAmelCase					:	Tuple      =       TER(
       normalized=_snake_case       ,  no_punct=_snake_case       ,  asian_support=_snake_case       ,  case_sensitive=_snake_case       ,  )
   __lowerCAmelCase					:	List[Any]      =       sb_ter.corpus_score(_snake_case       ,  _snake_case      )

   return {"score": output.score, "num_edits": output.num_edits, "ref_length": output.ref_length}