






"""simple docstring"""


from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block


@dataclass
class 					snake_case		(							_lowerCAmelCase						):


   '''simple docstring'''

   A_						:	torch.FloatTensor





class 					snake_case		(							nn.Module						):


   '''simple docstring'''

   def __init__(	self :     Optional[int],       _lowerCamelCase :     List[Any]=3,       _lowerCamelCase :     List[str]=3,       _lowerCamelCase :     Optional[int]=("DownEncoderBlock2D",),       _lowerCamelCase :     List[str]=(64,),       _lowerCamelCase :     int=2,       _lowerCamelCase :     str=32,       _lowerCamelCase :     Optional[Any]="silu",       _lowerCamelCase :     Optional[int]=True,       ):

     '''simple docstring'''
     super().__init__()
     __A									= layers_per_block

     __A									= torch.nn.Convad(
         _lowerCamelCase,       block_out_channels[0],       kernel_size=3,       stride=1,       padding=1,       )

     __A									= None
     __A									= nn.ModuleList([]  )

     # down
     __A									= block_out_channels[0]
     for i, down_block_type in enumerate(_lowerCamelCase  ):
       __A									= output_channel
       __A									= block_out_channels[i]
       __A									= i == len(_lowerCamelCase  ) - 1

       __A									= get_down_block(
           _lowerCamelCase,       num_layers=self.layers_per_block,       in_channels=_lowerCamelCase,       out_channels=_lowerCamelCase,       add_downsample=not is_final_block,       resnet_eps=1e-6,       downsample_padding=0,       resnet_act_fn=_lowerCamelCase,       resnet_groups=_lowerCamelCase,       attention_head_dim=_lowerCamelCase,       temb_channels=_lowerCamelCase,       )
       self.down_blocks.append(_lowerCamelCase  )

     # mid
     __A									= UNetMidBlockaD(
         in_channels=block_out_channels[-1],       resnet_eps=1e-6,       resnet_act_fn=_lowerCamelCase,       output_scale_factor=1,       resnet_time_scale_shift='''default''',       attention_head_dim=block_out_channels[-1],       resnet_groups=_lowerCamelCase,       temb_channels=_lowerCamelCase,       )

     # out
     __A									= nn.GroupNorm(num_channels=block_out_channels[-1],       num_groups=_lowerCamelCase,       eps=1e-6  )
     __A									= nn.SiLU()

     __A									= 2 * out_channels if double_z else out_channels
     __A									= nn.Convad(block_out_channels[-1],       _lowerCamelCase,       3,       padding=1  )

     __A									= False






   def 				_SCREAMING_SNAKE_CASE				(	self :     Union[str, Any],       _lowerCamelCase :     Dict  ):

     '''simple docstring'''
     __A									= x
     __A									= self.conv_in(_lowerCamelCase  )

     if self.training and self.gradient_checkpointing:

       def create_custom_forward(_lowerCamelCase :     Optional[Any]  ):
         def custom_forward(*_lowerCamelCase :     Optional[Any]  ):
           return module(*_lowerCamelCase  )

         return custom_forward

       # down
       if is_torch_version('''>=''',       '''1.11.0'''  ):
         for down_block in self.down_blocks:
           __A									= torch.utils.checkpoint.checkpoint(
               create_custom_forward(_lowerCamelCase  ),       _lowerCamelCase,       use_reentrant=_lowerCamelCase  )
         # middle
         __A									= torch.utils.checkpoint.checkpoint(
             create_custom_forward(self.mid_block  ),       _lowerCamelCase,       use_reentrant=_lowerCamelCase  )
       else:
         for down_block in self.down_blocks:
           __A									= torch.utils.checkpoint.checkpoint(create_custom_forward(_lowerCamelCase  ),       _lowerCamelCase  )
         # middle
         __A									= torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block  ),       _lowerCamelCase  )

     else:
       # down
       for down_block in self.down_blocks:
         __A									= down_block(_lowerCamelCase  )

       # middle
       __A									= self.mid_block(_lowerCamelCase  )

     # post-process
     __A									= self.conv_norm_out(_lowerCamelCase  )
     __A									= self.conv_act(_lowerCamelCase  )
     __A									= self.conv_out(_lowerCamelCase  )

     return sample





class 					snake_case		(							nn.Module						):


   '''simple docstring'''

   def __init__(	self :     Optional[Any],       _lowerCamelCase :     List[str]=3,       _lowerCamelCase :     List[Any]=3,       _lowerCamelCase :     Dict=("UpDecoderBlock2D",),       _lowerCamelCase :     List[str]=(64,),       _lowerCamelCase :     str=2,       _lowerCamelCase :     List[str]=32,       _lowerCamelCase :     Dict="silu",       _lowerCamelCase :     Optional[int]="group",       ):

     '''simple docstring'''
     super().__init__()
     __A									= layers_per_block

     __A									= nn.Convad(
         _lowerCamelCase,       block_out_channels[-1],       kernel_size=3,       stride=1,       padding=1,       )

     __A									= None
     __A									= nn.ModuleList([]  )

     __A									= in_channels if norm_type == '''spatial''' else None

     # mid
     __A									= UNetMidBlockaD(
         in_channels=block_out_channels[-1],       resnet_eps=1e-6,       resnet_act_fn=_lowerCamelCase,       output_scale_factor=1,       resnet_time_scale_shift='''default''' if norm_type == '''group''' else norm_type,       attention_head_dim=block_out_channels[-1],       resnet_groups=_lowerCamelCase,       temb_channels=_lowerCamelCase,       )

     # up
     __A									= list(reversed(_lowerCamelCase  )  )
     __A									= reversed_block_out_channels[0]
     for i, up_block_type in enumerate(_lowerCamelCase  ):
       __A									= output_channel
       __A									= reversed_block_out_channels[i]

       __A									= i == len(_lowerCamelCase  ) - 1

       __A									= get_up_block(
           _lowerCamelCase,       num_layers=self.layers_per_block + 1,       in_channels=_lowerCamelCase,       out_channels=_lowerCamelCase,       prev_output_channel=_lowerCamelCase,       add_upsample=not is_final_block,       resnet_eps=1e-6,       resnet_act_fn=_lowerCamelCase,       resnet_groups=_lowerCamelCase,       attention_head_dim=_lowerCamelCase,       temb_channels=_lowerCamelCase,       resnet_time_scale_shift=_lowerCamelCase,       )
       self.up_blocks.append(_lowerCamelCase  )
       __A									= output_channel

     # out
     if norm_type == "spatial":
       __A									= SpatialNorm(block_out_channels[0],       _lowerCamelCase  )
     else:
       __A									= nn.GroupNorm(num_channels=block_out_channels[0],       num_groups=_lowerCamelCase,       eps=1e-6  )
     __A									= nn.SiLU()
     __A									= nn.Convad(block_out_channels[0],       _lowerCamelCase,       3,       padding=1  )

     __A									= False






   def 				_SCREAMING_SNAKE_CASE				(	self :     List[Any],       _lowerCamelCase :     Optional[Any],       _lowerCamelCase :     Union[str, Any]=None  ):

     '''simple docstring'''
     __A									= z
     __A									= self.conv_in(_lowerCamelCase  )

     __A									= next(iter(self.up_blocks.parameters()  )  ).dtype
     if self.training and self.gradient_checkpointing:

       def create_custom_forward(_lowerCamelCase :     str  ):
         def custom_forward(*_lowerCamelCase :     Optional[int]  ):
           return module(*_lowerCamelCase  )

         return custom_forward

       if is_torch_version('''>=''',       '''1.11.0'''  ):
         # middle
         __A									= torch.utils.checkpoint.checkpoint(
             create_custom_forward(self.mid_block  ),       _lowerCamelCase,       _lowerCamelCase,       use_reentrant=_lowerCamelCase  )
         __A									= sample.to(_lowerCamelCase  )

         # up
         for up_block in self.up_blocks:
           __A									= torch.utils.checkpoint.checkpoint(
               create_custom_forward(_lowerCamelCase  ),       _lowerCamelCase,       _lowerCamelCase,       use_reentrant=_lowerCamelCase  )
       else:
         # middle
         __A									= torch.utils.checkpoint.checkpoint(
             create_custom_forward(self.mid_block  ),       _lowerCamelCase,       _lowerCamelCase  )
         __A									= sample.to(_lowerCamelCase  )

         # up
         for up_block in self.up_blocks:
           __A									= torch.utils.checkpoint.checkpoint(create_custom_forward(_lowerCamelCase  ),       _lowerCamelCase,       _lowerCamelCase  )
     else:
       # middle
       __A									= self.mid_block(_lowerCamelCase,       _lowerCamelCase  )
       __A									= sample.to(_lowerCamelCase  )

       # up
       for up_block in self.up_blocks:
         __A									= up_block(_lowerCamelCase,       _lowerCamelCase  )

        # post-process
     if latent_embeds is None:
       __A									= self.conv_norm_out(_lowerCamelCase  )
     else:
       __A									= self.conv_norm_out(_lowerCamelCase,       _lowerCamelCase  )
     __A									= self.conv_act(_lowerCamelCase  )
     __A									= self.conv_out(_lowerCamelCase  )

     return sample





class 					snake_case		(							nn.Module						):


   '''simple docstring'''

   def __init__(	self :     Union[str, Any],       _lowerCamelCase :     int,       _lowerCamelCase :     str,       _lowerCamelCase :     Tuple,       _lowerCamelCase :     Any=None,       _lowerCamelCase :     Tuple="random",       _lowerCamelCase :     Dict=False,       _lowerCamelCase :     Any=True  ):

     '''simple docstring'''
     super().__init__()
     __A									= n_e
     __A									= vq_embed_dim
     __A									= beta
     __A									= legacy

     __A									= nn.Embedding(self.n_e,       self.vq_embed_dim  )
     self.embedding.weight.data.uniform_(-1.0 / self.n_e,       1.0 / self.n_e  )

     __A									= remap
     if self.remap is not None:
       self.register_buffer('''used''',       torch.tensor(np.load(self.remap  )  )  )
       __A									= self.used.shape[0]
       __A									= unknown_index  # "random" or "extra" or integer
       if self.unknown_index == "extra":
         __A									= self.re_embed
         __A									= self.re_embed + 1
       print(
           f'Remapping {self.n_e} indices to {self.re_embed} indices. '
           f'Using {self.unknown_index} for unknown indices.'  )
     else:
       __A									= n_e

     __A									= sane_index_shape


   def 				_SCREAMING_SNAKE_CASE				(	self :     Dict,       _lowerCamelCase :     int  ):

     '''simple docstring'''
     __A									= inds.shape
     assert len(_lowerCamelCase  ) > 1
     __A									= inds.reshape(ishape[0],       -1  )
     __A									= self.used.to(_lowerCamelCase  )
     __A									= (inds[:, :, None] == used[None, None, ...]).long()
     __A									= match.argmax(-1  )
     __A									= match.sum(2  ) < 1
     if self.unknown_index == "random":
       __A									= torch.randint(0,       self.re_embed,       size=new[unknown].shape  ).to(device=new.device  )
     else:
       __A									= self.unknown_index
     return new.reshape(_lowerCamelCase  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Dict,       _lowerCamelCase :     Union[str, Any]  ):

     '''simple docstring'''
     __A									= inds.shape
     assert len(_lowerCamelCase  ) > 1
     __A									= inds.reshape(ishape[0],       -1  )
     __A									= self.used.to(_lowerCamelCase  )
     if self.re_embed > self.used.shape[0]:  # extra token
       __A									= 0  # simply set to zero
     __A									= torch.gather(used[None, :][inds.shape[0] * [0], :],       1,       _lowerCamelCase  )
     return back.reshape(_lowerCamelCase  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     List[str],       _lowerCamelCase :     Any  ):

     '''simple docstring'''
     # reshape z -> (batch, height, width, channel) and flatten
     __A									= z.permute(0,       2,       3,       1  ).contiguous()
     __A									= z.view(-1,       self.vq_embed_dim  )

     # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
     __A									= torch.argmin(torch.cdist(_lowerCamelCase,       self.embedding.weight  ),       dim=1  )

     __A									= self.embedding(_lowerCamelCase  ).view(z.shape  )
     __A									= None
     __A									= None

     # compute loss for embedding
     if not self.legacy:
       __A									= self.beta * torch.mean((z_q.detach() - z) ** 2  ) + torch.mean((z_q - z.detach()) ** 2  )
     else:
       __A									= torch.mean((z_q.detach() - z) ** 2  ) + self.beta * torch.mean((z_q - z.detach()) ** 2  )

     # preserve gradients
     __A									= z + (z_q - z).detach()

     # reshape back to match original input shape
     __A									= z_q.permute(0,       3,       1,       2  ).contiguous()

     if self.remap is not None:
       __A									= min_encoding_indices.reshape(z.shape[0],       -1  )  # add batch axis
       __A									= self.remap_to_used(_lowerCamelCase  )
       __A									= min_encoding_indices.reshape(-1,       1  )  # flatten

     if self.sane_index_shape:
       __A									= min_encoding_indices.reshape(z_q.shape[0],       z_q.shape[2],       z_q.shape[3]  )

     return z_q, loss, (perplexity, min_encodings, min_encoding_indices)






   def 				_SCREAMING_SNAKE_CASE				(	self :     Union[str, Any],       _lowerCamelCase :     Optional[Any],       _lowerCamelCase :     List[Any]  ):

     '''simple docstring'''
     # shape specifying (batch, height, width, channel)
     if self.remap is not None:
       __A									= indices.reshape(shape[0],       -1  )  # add batch axis
       __A									= self.unmap_to_all(_lowerCamelCase  )
       __A									= indices.reshape(-1  )  # flatten again

     # get quantized latent vectors
     __A									= self.embedding(_lowerCamelCase  )

     if shape is not None:
       __A									= z_q.view(_lowerCamelCase  )
       # reshape back to match original input shape
       __A									= z_q.permute(0,       3,       1,       2  ).contiguous()

     return z_q





class 					snake_case		(							_lowerCAmelCase						):


   '''simple docstring'''

   def __init__(	self :     Any,       _lowerCamelCase :     str,       _lowerCamelCase :     List[Any]=False  ):

     '''simple docstring'''
     __A									= parameters
     __A			,    __A									= torch.chunk(_lowerCamelCase,       2,       dim=1  )
     __A									= torch.clamp(self.logvar,       -30.0,       20.0  )
     __A									= deterministic
     __A									= torch.exp(0.5 * self.logvar  )
     __A									= torch.exp(self.logvar  )
     if self.deterministic:
       __A									= __A									= torch.zeros_like(
           self.mean,       device=self.parameters.device,       dtype=self.parameters.dtype  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[int],       _lowerCamelCase :     Optional[torch.Generator] = None  ):

     '''simple docstring'''
     # make sure sample is on the same device as the parameters and has same dtype
     __A									= randn_tensor(
         self.mean.shape,       generator=_lowerCamelCase,       device=self.parameters.device,       dtype=self.parameters.dtype  )
     __A									= self.mean + self.std * sample
     return x


   def 				_SCREAMING_SNAKE_CASE				(	self :     Any,       _lowerCamelCase :     Optional[int]=None  ):

     '''simple docstring'''
     if self.deterministic:
       return torch.Tensor([0.0]  )
     else:
       if other is None:
         return 0.5 * torch.sum(torch.pow(self.mean,       2  ) + self.var - 1.0 - self.logvar,       dim=[1, 2, 3]  )
       else:
         return 0.5 * torch.sum(
             torch.pow(self.mean - other.mean,       2  ) / other.var
             + self.var / other.var
             - 1.0
             - self.logvar
             + other.logvar,       dim=[1, 2, 3],       )


   def 				_SCREAMING_SNAKE_CASE				(	self :     int,       _lowerCamelCase :     Optional[Any],       _lowerCamelCase :     Tuple=[1, 2, 3]  ):

     '''simple docstring'''
     if self.deterministic:
       return torch.Tensor([0.0]  )
     __A									= np.log(2.0 * np.pi  )
     return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean,       2  ) / self.var,       dim=_lowerCamelCase  )






   def 				_SCREAMING_SNAKE_CASE				(	self :     Tuple  ):

     '''simple docstring'''
     return self.mean












"""simple docstring"""


import faiss  # noqa: F401 # Here to have a nice missing dependency error message early on
import numpy  # noqa: F401 # Here to have a nice missing dependency error message early on
import requests  # noqa: F401 # Here to have a nice missing dependency error message early on
import sklearn  # noqa: F401 # Here to have a nice missing dependency error message early on
import tqdm  # noqa: F401 # Here to have a nice missing dependency error message early on
from mauve import compute_mauve  # From: mauve-text

import datasets


lowercase_				     =     '\\n@inproceedings{pillutla-etal:mauve:neurips2021,\n  title={MAUVE: Measuring the Gap Between Neural Text and Human Text using Divergence Frontiers},\n  author={Pillutla, Krishna and Swayamdipta, Swabha and Zellers, Rowan and Thickstun, John and Welleck, Sean and Choi, Yejin and Harchaoui, Zaid},\n  booktitle = {NeurIPS},\n  year      = {2021}\n}\n\n'

lowercase_				     =     '\\nMAUVE is a library built on PyTorch and HuggingFace Transformers to measure the gap between neural text and human text with the eponymous MAUVE measure.\n\nMAUVE summarizes both Type I and Type II errors measured softly using Kullback–Leibler (KL) divergences.\n\nFor details, see the MAUVE paper: https://arxiv.org/abs/2102.01454 (Neurips, 2021).\n\nThis metrics is a wrapper around the official implementation of MAUVE:\nhttps://github.com/krishnap25/mauve\n'

lowercase_				     =     '\nCalculates MAUVE scores between two lists of generated text and reference text.\nArgs:\n    predictions: list of generated text to score. Each predictions\n        should be a string with tokens separated by spaces.\n    references: list of reference for each prediction. Each\n        reference should be a string with tokens separated by spaces.\nOptional Args:\n    num_buckets: the size of the histogram to quantize P and Q. Options: \'auto\' (default) or an integer\n    pca_max_data: the number data points to use for PCA dimensionality reduction prior to clustering. If -1, use all the data. Default -1\n    kmeans_explained_var: amount of variance of the data to keep in dimensionality reduction by PCA. Default 0.9\n    kmeans_num_redo: number of times to redo k-means clustering (the best objective is kept). Default 5\n    kmeans_max_iter: maximum number of k-means iterations. Default 500\n    featurize_model_name: name of the model from which features are obtained. Default \'gpt2-large\' Use one of [\'gpt2\', \'gpt2-medium\', \'gpt2-large\', \'gpt2-xl\'].\n    device_id: Device for featurization. Supply a GPU id (e.g. 0 or 3) to use GPU. If no GPU with this id is found, use CPU\n    max_text_length: maximum number of tokens to consider. Default 1024\n    divergence_curve_discretization_size: Number of points to consider on the divergence curve. Default 25\n    mauve_scaling_factor: "c" from the paper. Default 5.\n    verbose: If True (default), print running time updates\n    seed: random seed to initialize k-means cluster assignments.\nReturns:\n    mauve: MAUVE score, a number between 0 and 1. Larger values indicate that P and Q are closer,\n    frontier_integral: Frontier Integral, a number between 0 and 1. Smaller values indicate that P and Q are closer,\n    divergence_curve: a numpy.ndarray of shape (m, 2); plot it with matplotlib to view the divergence curve,\n    p_hist: a discrete distribution, which is a quantized version of the text distribution p_text,\n    q_hist: same as above, but with q_text.\nExamples:\n\n    >>> # faiss segfaults in doctest for some reason, so the .compute call is not tested with doctest\n    >>> import datasets\n    >>> mauve = datasets.load_metric(\'mauve\')\n    >>> predictions = ["hello there", "general kenobi"]\n    >>> references = ["hello there", "general kenobi"]\n    >>> out = mauve.compute(predictions=predictions, references=references) # doctest: +SKIP\n    >>> print(out.mauve) # doctest: +SKIP\n    1.0\n'


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


   '''simple docstring'''

   def 				_SCREAMING_SNAKE_CASE				(	self :     Any  ):

     '''simple docstring'''
     return datasets.MetricInfo(
         description=_DESCRIPTION,       citation=_CITATION,       homepage='''https://github.com/krishnap25/mauve''',       inputs_description=_KWARGS_DESCRIPTION,       features=datasets.Features(
             {
                 '''predictions''': datasets.Value('''string''',       id='''sequence'''  ),
                 '''references''': datasets.Value('''string''',       id='''sequence'''  ),
             }  ),       codebase_urls=['''https://github.com/krishnap25/mauve'''],       reference_urls=[
             '''https://arxiv.org/abs/2102.01454''',
             '''https://github.com/krishnap25/mauve''',
         ],       )






   def 				_SCREAMING_SNAKE_CASE				(	self :     int,       _lowerCamelCase :     str,       _lowerCamelCase :     Optional[Any],       _lowerCamelCase :     Any=None,       _lowerCamelCase :     Tuple=None,       _lowerCamelCase :     Optional[Any]=None,       _lowerCamelCase :     Union[str, Any]=None,       _lowerCamelCase :     str="auto",       _lowerCamelCase :     Union[str, Any]=-1,       _lowerCamelCase :     List[str]=0.9,       _lowerCamelCase :     int=5,       _lowerCamelCase :     Tuple=5_00,       _lowerCamelCase :     Union[str, Any]="gpt2-large",       _lowerCamelCase :     int=-1,       _lowerCamelCase :     Union[str, Any]=10_24,       _lowerCamelCase :     Union[str, Any]=25,       _lowerCamelCase :     str=5,       _lowerCamelCase :     Any=True,       _lowerCamelCase :     Union[str, Any]=25,       ):

     '''simple docstring'''
     __A									= compute_mauve(
         p_text=_lowerCamelCase,       q_text=_lowerCamelCase,       p_features=_lowerCamelCase,       q_features=_lowerCamelCase,       p_tokens=_lowerCamelCase,       q_tokens=_lowerCamelCase,       num_buckets=_lowerCamelCase,       pca_max_data=_lowerCamelCase,       kmeans_explained_var=_lowerCamelCase,       kmeans_num_redo=_lowerCamelCase,       kmeans_max_iter=_lowerCamelCase,       featurize_model_name=_lowerCamelCase,       device_id=_lowerCamelCase,       max_text_length=_lowerCamelCase,       divergence_curve_discretization_size=_lowerCamelCase,       mauve_scaling_factor=_lowerCamelCase,       verbose=_lowerCamelCase,       seed=_lowerCamelCase,       )
     return out




