



"""simple docstring"""
from jiwer import compute_measures

import datasets


__a      =				"\\n@inproceedings{inproceedings,\n    author = {Morris, Andrew and Maier, Viktoria and Green, Phil},\n    year = {2004},\n    month = {01},\n    pages = {},\n    title = {From WER and RIL to MER and WIL: improved evaluation measures for connected speech recognition.}\n}\n"

__a      =				"\\nWord error rate (WER) is a common metric of the performance of an automatic speech recognition system.\n\nThe general difficulty of measuring performance lies in the fact that the recognized word sequence can have a different length from the reference word sequence (supposedly the correct one). The WER is derived from the Levenshtein distance, working at the word level instead of the phoneme level. The WER is a valuable tool for comparing different systems as well as for evaluating improvements within one system. This kind of measurement, however, provides no details on the nature of translation errors and further work is therefore required to identify the main source(s) of error and to focus any research effort.\n\nThis problem is solved by first aligning the recognized word sequence with the reference (spoken) word sequence using dynamic string alignment. Examination of this issue is seen through a theory called the power law that states the correlation between perplexity and word error rate.\n\nWord error rate can then be computed as:\n\nWER = (S + D + I) / N = (S + D + I) / (S + D + C)\n\nwhere\n\nS is the number of substitutions,\nD is the number of deletions,\nI is the number of insertions,\nC is the number of correct words,\nN is the number of words in the reference (N=S+D+C).\n\nThis value indicates the average number of errors per reference word. The lower the value, the better the\nperformance of the ASR system with a WER of 0 being a perfect score.\n"

__a      =				"\nCompute WER score of transcribed segments against references.\n\nArgs:\n    references: List of references for each speech input.\n    predictions: List of transcriptions to score.\n    concatenate_texts (bool, default=False): Whether to concatenate all input texts or compute WER iteratively.\n\nReturns:\n    (float): the word error rate\n\nExamples:\n\n    >>> predictions = [\"this is the prediction\", \"there is an other sample\"]\n    >>> references = [\"this is the reference\", \"there is another one\"]\n    >>> wer = datasets.load_metric(\"wer\")\n    >>> wer_score = wer.compute(predictions=predictions, references=references)\n    >>> print(wer_score)\n    0.5\n"
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION    ,						_KWARGS_DESCRIPTION				)
class 				lowerCamelCase  (			datasets.Metric				):






					'''simple docstring'''





					def    lowerCAmelCase_				(       self:							Optional[Any]						)   ->	Optional[Any]:
											return datasets.MetricInfo(
											    description=_DESCRIPTION						,     citation=_CITATION						,     inputs_description=_KWARGS_DESCRIPTION						,     features=datasets.Features(
											        {
											            """predictions""": datasets.Value("""string"""						,     id="""sequence"""						),
											            """references""": datasets.Value("""string"""						,     id="""sequence"""						),
											        }						)						,     codebase_urls=["""https://github.com/jitsi/jiwer/"""]						,     reference_urls=[
											        """https://en.wikipedia.org/wiki/Word_error_rate""",
											    ]						,     )






					def    lowerCAmelCase_				(       self:							int						,     snake_case:							Optional[Any]=None						,     snake_case:							Dict=None						,     snake_case:							Any=False						)   ->	Optional[int]:
											if concatenate_texts:
																	return compute_measures(snake_case						,     snake_case						)["wer"]
											else:
																	snake_case_ :List[str]						=			0
																	snake_case_ :Dict						=			0
																	for prediction, reference in zip(snake_case						,     snake_case						):
																							snake_case_ :List[str]						=			compute_measures(snake_case						,     snake_case						)
																							incorrect += measures["substitutions"] + measures["deletions"] + measures["insertions"]
																							total += measures["substitutions"] + measures["deletions"] + measures["hits"]
																	return incorrect / total

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







@flax.struct.dataclass
class 				_lowercase     (				UpperCAmelCase__			):


 '''simple docstring'''






 _SCREAMING_SNAKE_CASE     :       jnp.ndarray







@flax_register_to_config
class 				_lowercase     (				nn.Module  ,							UpperCAmelCase__  ,							UpperCAmelCase__			):


 '''simple docstring'''






 _SCREAMING_SNAKE_CASE     :       int      	=					32
 _SCREAMING_SNAKE_CASE     :       int      	=					4
 _SCREAMING_SNAKE_CASE     :       int      	=					4
 _SCREAMING_SNAKE_CASE     :       Tuple[str]      	=					(
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "DownBlock2D",
 )
 _SCREAMING_SNAKE_CASE     :       Tuple[str]      	=					("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
 _SCREAMING_SNAKE_CASE     :       Union[bool, Tuple[bool]]      	=					False
 _SCREAMING_SNAKE_CASE     :       Tuple[int]      	=					(320, 640, 1280, 1280)
 _SCREAMING_SNAKE_CASE     :       int      	=					2
 _SCREAMING_SNAKE_CASE     :       Union[int, Tuple[int]]      	=					8
 _SCREAMING_SNAKE_CASE     :       Optional[Union[int, Tuple[int]]]      	=					None
 _SCREAMING_SNAKE_CASE     :       int      	=					1280
 _SCREAMING_SNAKE_CASE     :       float      	=					0.0
 _SCREAMING_SNAKE_CASE     :       bool      	=					False
 _SCREAMING_SNAKE_CASE     :       jnp.dtype      	=					jnp.floataa
 _SCREAMING_SNAKE_CASE     :       bool      	=					True
 _SCREAMING_SNAKE_CASE     :       int      	=					0
 _SCREAMING_SNAKE_CASE     :       bool      	=					False
 def a					(      self   :	Dict      ,       SCREAMING_SNAKE_CASE__   :	jax.random.KeyArray							)  ->    FrozenDict:
  # init input tensors
  __lowerCAmelCase											=						(1, self.in_channels, self.sample_size, self.sample_size)
  __lowerCAmelCase											=						jnp.zeros(SCREAMING_SNAKE_CASE__      ,       dtype=jnp.floataa							)
  __lowerCAmelCase											=						jnp.ones((1,)      ,       dtype=jnp.intaa							)
  __lowerCAmelCase											=						jnp.zeros((1, 1, self.cross_attention_dim)      ,       dtype=jnp.floataa							)

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

  return self.init(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__							)["params"]
 def a					(      self   :	int							)  ->    List[str]:
  __lowerCAmelCase											=						self.block_out_channels
  __lowerCAmelCase											=						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											=						self.num_attention_heads or self.attention_head_dim

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

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

  __lowerCAmelCase											=						self.only_cross_attention
  if isinstance(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__							):
   __lowerCAmelCase											=						(only_cross_attention,) * len(self.down_block_types							)

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

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

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

   down_blocks.append(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						down_blocks

  # mid
  __lowerCAmelCase											=						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											=						[]
  __lowerCAmelCase											=						list(reversed(SCREAMING_SNAKE_CASE__							)							)
  __lowerCAmelCase											=						list(reversed(SCREAMING_SNAKE_CASE__							)							)
  __lowerCAmelCase											=						list(reversed(SCREAMING_SNAKE_CASE__							)							)
  __lowerCAmelCase											=						reversed_block_out_channels[0]
  for i, up_block_type in enumerate(self.up_block_types							):
   __lowerCAmelCase											=						output_channel
   __lowerCAmelCase											=						reversed_block_out_channels[i]
   __lowerCAmelCase											=						reversed_block_out_channels[min(i + 1      ,       len(SCREAMING_SNAKE_CASE__							) - 1							)]

   __lowerCAmelCase											=						i == len(SCREAMING_SNAKE_CASE__							) - 1

   if up_block_type == "CrossAttnUpBlock2D":
    __lowerCAmelCase											=						FlaxCrossAttnUpBlockaD(
        in_channels=SCREAMING_SNAKE_CASE__      ,       out_channels=SCREAMING_SNAKE_CASE__      ,       prev_output_channel=SCREAMING_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											=						FlaxUpBlockaD(
        in_channels=SCREAMING_SNAKE_CASE__      ,       out_channels=SCREAMING_SNAKE_CASE__      ,       prev_output_channel=SCREAMING_SNAKE_CASE__      ,       num_layers=self.layers_per_block + 1      ,       add_upsample=not is_final_block      ,       dropout=self.dropout      ,       dtype=self.dtype      ,       )

   up_blocks.append(SCREAMING_SNAKE_CASE__							)
   __lowerCAmelCase											=						output_channel
  __lowerCAmelCase											=						up_blocks

  # out
  __lowerCAmelCase											=						nn.GroupNorm(num_groups=32      ,       epsilon=1e-5							)
  __lowerCAmelCase											=						nn.Conv(
      self.out_channels      ,       kernel_size=(3, 3)      ,       strides=(1, 1)      ,       padding=((1, 1), (1, 1))      ,       dtype=self.dtype      ,       )





 def __call__(      self   :	Dict      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	Tuple      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]      ,       SCREAMING_SNAKE_CASE__   :	Optional[Any]=None      ,       SCREAMING_SNAKE_CASE__   :	List[str]=None      ,       SCREAMING_SNAKE_CASE__   :	bool = True      ,       SCREAMING_SNAKE_CASE__   :	bool = False      ,       )  ->    Union[FlaxUNetaDConditionOutput, Tuple]:
  # 1. time
  if not isinstance(SCREAMING_SNAKE_CASE__      ,       jnp.ndarray							):
   __lowerCAmelCase											=						jnp.array([timesteps]      ,       dtype=jnp.intaa							)
  elif isinstance(SCREAMING_SNAKE_CASE__      ,       jnp.ndarray							) and len(timesteps.shape							) == 0:
   __lowerCAmelCase											=						timesteps.astype(dtype=jnp.floataa							)
   __lowerCAmelCase											=						jnp.expand_dims(SCREAMING_SNAKE_CASE__      ,       0							)

  __lowerCAmelCase											=						self.time_proj(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						self.time_embedding(SCREAMING_SNAKE_CASE__							)

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

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

  if down_block_additional_residuals is not None:
   __lowerCAmelCase											=						()

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

   __lowerCAmelCase											=						new_down_block_res_samples

  # 4. mid
  __lowerCAmelCase											=						self.mid_block(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       SCREAMING_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											=						down_block_res_samples[-(self.layers_per_block + 1) :]
   __lowerCAmelCase											=						down_block_res_samples[: -(self.layers_per_block + 1)]
   if isinstance(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__							):
    __lowerCAmelCase											=						up_block(
        SCREAMING_SNAKE_CASE__      ,       temb=SCREAMING_SNAKE_CASE__      ,       encoder_hidden_states=SCREAMING_SNAKE_CASE__      ,       res_hidden_states_tuple=SCREAMING_SNAKE_CASE__      ,       deterministic=not train      ,       )
   else:
    __lowerCAmelCase											=						up_block(SCREAMING_SNAKE_CASE__      ,       temb=SCREAMING_SNAKE_CASE__      ,       res_hidden_states_tuple=SCREAMING_SNAKE_CASE__      ,       deterministic=not train							)

        # 6. post-process
  __lowerCAmelCase											=						self.conv_norm_out(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						nn.silu(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						self.conv_out(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						jnp.transpose(SCREAMING_SNAKE_CASE__      ,       (0, 3, 1, 2)							)

  if not return_dict:
   return (sample,)

  return FlaxUNetaDConditionOutput(sample=SCREAMING_SNAKE_CASE__							)






