


"""simple docstring"""
import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEmbeddings,
    BertLayer,
    BertPooler,
    BertPreTrainedModel,
)







def   _a			( _SCREAMING_SNAKE_CASE			)      ->			Dict:
       snake_case_          =						torch.exp(_SCREAMING_SNAKE_CASE			)
       snake_case_          =						torch.sum(_SCREAMING_SNAKE_CASE			,	dim=1			)  # sum of exp(x_i)
       snake_case_          =						torch.sum(x * exp_x			,	dim=1			)  # sum of x_i * exp(x_i)
       return torch.log(_SCREAMING_SNAKE_CASE			) - B / A





class        __A       (nn.Module):





       '''simple docstring'''

       def __init__(			self							: Optional[Any]						,    UpperCAmelCase_							: Tuple    )   ->int:



              """simple docstring"""



              super().__init__()
              snake_case_          =						config.output_attentions
              snake_case_          =						config.output_hidden_states
              snake_case_          =						nn.ModuleList([BertLayer(UpperCAmelCase_    ) for _ in range(config.num_hidden_layers    )]    )
              snake_case_          =						nn.ModuleList([BertHighway(UpperCAmelCase_    ) for _ in range(config.num_hidden_layers    )]    )

              snake_case_          =						[-1 for _ in range(config.num_hidden_layers    )]



       def 						lowerCAmelCase    (			self							: str						,    UpperCAmelCase_							: int    )   ->Optional[int]:



              """simple docstring"""



              if (type(UpperCAmelCase_    ) is float) or (type(UpperCAmelCase_    ) is int):
                     for i in range(len(self.early_exit_entropy    )    ):
                            snake_case_          =						x
              else:
                     snake_case_          =						x



       def 						lowerCAmelCase    (			self							: Dict						,    UpperCAmelCase_							: List[str]    )   ->Optional[Any]:



              """simple docstring"""



              snake_case_          =						pooler.state_dict()
              for highway in self.highway:
                     for name, param in highway.pooler.state_dict().items():
                            param.copy_(loaded_model[name]    )



       def 						lowerCAmelCase    (			self							: List[Any]						,    UpperCAmelCase_							: List[Any]						,    UpperCAmelCase_							: Optional[Any]=None						,    UpperCAmelCase_							: List[str]=None						,    UpperCAmelCase_							: Optional[Any]=None						,    UpperCAmelCase_							: Any=None						,    )   ->Optional[Any]:



              """simple docstring"""



              snake_case_          =						()
              snake_case_          =						()
              snake_case_          =						()
              for i, layer_module in enumerate(self.layer    ):
                     if self.output_hidden_states:
                            snake_case_          =						all_hidden_states + (hidden_states,)

                     snake_case_          =						layer_module(
                         UpperCAmelCase_						,    UpperCAmelCase_						,    head_mask[i]						,    UpperCAmelCase_						,    UpperCAmelCase_    )
                     snake_case_          =						layer_outputs[0]

                     if self.output_attentions:
                            snake_case_          =						all_attentions + (layer_outputs[1],)

                     snake_case_          =						(hidden_states,)
                     if self.output_hidden_states:
                            snake_case_          =						current_outputs + (all_hidden_states,)
                     if self.output_attentions:
                            snake_case_          =						current_outputs + (all_attentions,)

                     snake_case_          =						self.highway[i](UpperCAmelCase_    )
                     # logits, pooled_output

                     if not self.training:
                            snake_case_          =						highway_exit[0]
                            snake_case_          =						entropy(UpperCAmelCase_    )
                            snake_case_          =						highway_exit + (highway_entropy,)  # logits, hidden_states(?), entropy
                            snake_case_          =						all_highway_exits + (highway_exit,)

                            if highway_entropy < self.early_exit_entropy[i]:
                                   snake_case_          =						(highway_logits,) + current_outputs[1:] + (all_highway_exits,)
                                   raise HighwayException(UpperCAmelCase_						,    i + 1    )
                     else:
                            snake_case_          =						all_highway_exits + (highway_exit,)

        # Add last layer
              if self.output_hidden_states:
                     snake_case_          =						all_hidden_states + (hidden_states,)

              snake_case_          =						(hidden_states,)
              if self.output_hidden_states:
                     snake_case_          =						outputs + (all_hidden_states,)
              if self.output_attentions:
                     snake_case_          =						outputs + (all_attentions,)

              snake_case_          =						outputs + (all_highway_exits,)
              return outputs  # last-layer hidden state, (all hidden states), (all attentions), all highway exits





@add_start_docstrings(
    """The Bert Model transformer with early exiting (DeeBERT). """  ,						snake_case__  ,						)
class        __A       (snake_case__):





       '''simple docstring'''

       def __init__(			self							: Any						,    UpperCAmelCase_							: List[str]    )   ->Tuple:



              """simple docstring"""



              super().__init__(UpperCAmelCase_    )
              snake_case_          =						config

              snake_case_          =						BertEmbeddings(UpperCAmelCase_    )
              snake_case_          =						DeeBertEncoder(UpperCAmelCase_    )
              snake_case_          =						BertPooler(UpperCAmelCase_    )

              self.init_weights()



       def 						lowerCAmelCase    (			self							: Union[str, Any]    )   ->Optional[Any]:



              """simple docstring"""



              self.encoder.init_highway_pooler(self.pooler    )



       def 						lowerCAmelCase    (			self							: str    )   ->Union[str, Any]:



              """simple docstring"""



              return self.embeddings.word_embeddings



       def 						lowerCAmelCase    (			self							: Optional[Any]						,    UpperCAmelCase_							: Union[str, Any]    )   ->Any:



              """simple docstring"""



              snake_case_          =						value



       def 						lowerCAmelCase    (			self							: Any						,    UpperCAmelCase_							: Tuple    )   ->Optional[int]:



              """simple docstring"""



              for layer, heads in heads_to_prune.items():
                     self.encoder.layer[layer].attention.prune_heads(UpperCAmelCase_    )



       @add_start_docstrings_to_model_forward(UpperCAmelCase_    )
       def 						lowerCAmelCase    (			self							: List[Any]						,    UpperCAmelCase_							: str=None						,    UpperCAmelCase_							: Any=None						,    UpperCAmelCase_							: int=None						,    UpperCAmelCase_							: Tuple=None						,    UpperCAmelCase_							: int=None						,    UpperCAmelCase_							: Tuple=None						,    UpperCAmelCase_							: str=None						,    UpperCAmelCase_							: int=None						,    )   ->Optional[int]:



              """simple docstring"""



              if input_ids is not None and inputs_embeds is not None:
                     raise ValueError("""You cannot specify both input_ids and inputs_embeds at the same time"""    )
              elif input_ids is not None:
                     snake_case_          =						input_ids.size()
              elif inputs_embeds is not None:
                     snake_case_          =						inputs_embeds.size()[:-1]
              else:
                     raise ValueError("""You have to specify either input_ids or inputs_embeds"""    )

              snake_case_          =						input_ids.device if input_ids is not None else inputs_embeds.device

              if attention_mask is None:
                     snake_case_          =						torch.ones(UpperCAmelCase_						,    device=UpperCAmelCase_    )
              if encoder_attention_mask is None:
                     snake_case_          =						torch.ones(UpperCAmelCase_						,    device=UpperCAmelCase_    )
              if token_type_ids is None:
                     snake_case_          =						torch.zeros(UpperCAmelCase_						,    dtype=torch.long						,    device=UpperCAmelCase_    )

              # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
              # ourselves in which case we just need to make it broadcastable to all heads.
              snake_case_          =						self.get_extended_attention_mask(UpperCAmelCase_						,    UpperCAmelCase_						,    UpperCAmelCase_    )

              # If a 2D ou 3D attention mask is provided for the cross-attention
              # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
              if encoder_attention_mask.dim() == 3:
                     snake_case_          =						encoder_attention_mask[:, None, :, :]
              if encoder_attention_mask.dim() == 2:
                     snake_case_          =						encoder_attention_mask[:, None, None, :]

              snake_case_          =						encoder_extended_attention_mask.to(
                  dtype=next(self.parameters()    ).dtype    )  # fp16 compatibility
              snake_case_          =						(1.0 - encoder_extended_attention_mask) * -10_000.0

              # Prepare head mask if needed
              # 1.0 in head_mask indicate we keep the head
              # attention_probs has shape bsz x n_heads x N x N
              # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
              # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
              snake_case_          =						self.get_head_mask(UpperCAmelCase_						,    self.config.num_hidden_layers    )

              snake_case_          =						self.embeddings(
                  input_ids=UpperCAmelCase_						,    position_ids=UpperCAmelCase_						,    token_type_ids=UpperCAmelCase_						,    inputs_embeds=UpperCAmelCase_    )
              snake_case_          =						self.encoder(
                  UpperCAmelCase_						,    attention_mask=UpperCAmelCase_						,    head_mask=UpperCAmelCase_						,    encoder_hidden_states=UpperCAmelCase_						,    encoder_attention_mask=UpperCAmelCase_						,    )
              snake_case_          =						encoder_outputs[0]
              snake_case_          =						self.pooler(UpperCAmelCase_    )

              snake_case_          =						(
                  sequence_output,
                  pooled_output,
              ) + encoder_outputs[
                  1:
              ]  # add hidden_states and attentions if they are here
              return outputs  # sequence_output, pooled_output, (hidden_states), (attentions), highway exits





class        __A       (snake_case__):





       '''simple docstring'''

       def __init__(			self							: Any						,    UpperCAmelCase_							: int						,    UpperCAmelCase_							: Optional[Any]    )   ->Union[str, Any]:



              """simple docstring"""



              snake_case_          =						message
              snake_case_          =						exit_layer  # start from 1!





class        __A       (nn.Module):





       '''simple docstring'''

       def __init__(			self							: Any						,    UpperCAmelCase_							: str    )   ->List[Any]:



              """simple docstring"""



              super().__init__()
              snake_case_          =						BertPooler(UpperCAmelCase_    )
              snake_case_          =						nn.Dropout(config.hidden_dropout_prob    )
              snake_case_          =						nn.Linear(config.hidden_size						,    config.num_labels    )



       def 						lowerCAmelCase    (			self							: Dict						,    UpperCAmelCase_							: Dict    )   ->List[Any]:



              """simple docstring"""



              snake_case_          =						encoder_outputs[0]
              snake_case_          =						self.pooler(UpperCAmelCase_    )
              # "return" pooler_output

              # BertModel
              snake_case_          =						(pooler_input, pooler_output) + encoder_outputs[1:]
              # "return" bmodel_output

              # Dropout and classification
              snake_case_          =						bmodel_output[1]

              snake_case_          =						self.dropout(UpperCAmelCase_    )
              snake_case_          =						self.classifier(UpperCAmelCase_    )

              return logits, pooled_output





@add_start_docstrings(
    """Bert Model (with early exiting - DeeBERT) with a classifier on top,
    also takes care of multi-layer training. """  ,						snake_case__  ,						)
class        __A       (snake_case__):





       '''simple docstring'''

       def __init__(			self							: Tuple						,    UpperCAmelCase_							: Optional[int]    )   ->Union[str, Any]:



              """simple docstring"""



              super().__init__(UpperCAmelCase_    )
              snake_case_          =						config.num_labels
              snake_case_          =						config.num_hidden_layers

              snake_case_          =						DeeBertModel(UpperCAmelCase_    )
              snake_case_          =						nn.Dropout(config.hidden_dropout_prob    )
              snake_case_          =						nn.Linear(config.hidden_size						,    self.config.num_labels    )

              self.init_weights()



       @add_start_docstrings_to_model_forward(UpperCAmelCase_    )
       def 						lowerCAmelCase    (			self							: List[str]						,    UpperCAmelCase_							: Optional[Any]=None						,    UpperCAmelCase_							: Any=None						,    UpperCAmelCase_							: Dict=None						,    UpperCAmelCase_							: Any=None						,    UpperCAmelCase_							: str=None						,    UpperCAmelCase_							: Any=None						,    UpperCAmelCase_							: Dict=None						,    UpperCAmelCase_							: Optional[Any]=-1						,    UpperCAmelCase_							: Optional[Any]=False						,    )   ->List[Any]:



              """simple docstring"""



              snake_case_          =						self.num_layers
              try:
                     snake_case_          =						self.bert(
                         UpperCAmelCase_						,    attention_mask=UpperCAmelCase_						,    token_type_ids=UpperCAmelCase_						,    position_ids=UpperCAmelCase_						,    head_mask=UpperCAmelCase_						,    inputs_embeds=UpperCAmelCase_						,    )
                     # sequence_output, pooled_output, (hidden_states), (attentions), highway exits

                     snake_case_          =						outputs[1]

                     snake_case_          =						self.dropout(UpperCAmelCase_    )
                     snake_case_          =						self.classifier(UpperCAmelCase_    )
                     snake_case_          =						(logits,) + outputs[2:]  # add hidden states and attention if they are here
              except HighwayException as e:
                     snake_case_          =						e.message
                     snake_case_          =						e.exit_layer
                     snake_case_          =						outputs[0]

              if not self.training:
                     snake_case_          =						entropy(UpperCAmelCase_    )
                     snake_case_          =						[]
                     snake_case_          =						[]
              if labels is not None:
                     if self.num_labels == 1:
                            #  We are doing regression
                            snake_case_          =						MSELoss()
                            snake_case_          =						loss_fct(logits.view(-1    )						,    labels.view(-1    )    )
                     else:
                            snake_case_          =						CrossEntropyLoss()
                            snake_case_          =						loss_fct(logits.view(-1						,    self.num_labels    )						,    labels.view(-1    )    )

                     # work with highway exits
                     snake_case_          =						[]
                     for highway_exit in outputs[-1]:
                            snake_case_          =						highway_exit[0]
                            if not self.training:
                                   highway_logits_all.append(UpperCAmelCase_    )
                                   highway_entropy.append(highway_exit[2]    )
                            if self.num_labels == 1:
                                   #  We are doing regression
                                   snake_case_          =						MSELoss()
                                   snake_case_          =						loss_fct(highway_logits.view(-1    )						,    labels.view(-1    )    )
                            else:
                                   snake_case_          =						CrossEntropyLoss()
                                   snake_case_          =						loss_fct(highway_logits.view(-1						,    self.num_labels    )						,    labels.view(-1    )    )
                            highway_losses.append(UpperCAmelCase_    )

                     if train_highway:
                            snake_case_          =						(sum(highway_losses[:-1]    ),) + outputs
                            # exclude the final highway, of course
                     else:
                            snake_case_          =						(loss,) + outputs
              if not self.training:
                     snake_case_          =						outputs + ((original_entropy, highway_entropy), exit_layer)
                     if output_layer >= 0:
                            snake_case_          =						(
                                (outputs[0],) + (highway_logits_all[output_layer],) + outputs[2:]
                            )  # use the highway of the last layer

              return outputs  # (loss), logits, (hidden_states), (attentions), (highway_exits)







"""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


__SCREAMING_SNAKE_CASE							:      Dict						=      '\\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'

__SCREAMING_SNAKE_CASE							:      Optional[int]						=      '\\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'

__SCREAMING_SNAKE_CASE							:      Tuple						=      '\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        __A       (datasets.Metric):





							'''simple docstring'''

							def 						lowerCAmelCase    (			self							: Tuple    )   ->Union[str, 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 						lowerCAmelCase    (			self							: Optional[int]						,    UpperCAmelCase_							: Tuple						,    UpperCAmelCase_							: Any						,    UpperCAmelCase_							: Optional[Any]=None						,    UpperCAmelCase_							: Optional[Any]=None						,    UpperCAmelCase_							: List[Any]=None						,    UpperCAmelCase_							: Any=None						,    UpperCAmelCase_							: int="auto"						,    UpperCAmelCase_							: Dict=-1						,    UpperCAmelCase_							: Optional[Any]=0.9						,    UpperCAmelCase_							: Dict=5						,    UpperCAmelCase_							: Optional[int]=500						,    UpperCAmelCase_							: Any="gpt2-large"						,    UpperCAmelCase_							: Union[str, Any]=-1						,    UpperCAmelCase_							: Optional[Any]=1_024						,    UpperCAmelCase_							: Dict=25						,    UpperCAmelCase_							: Optional[Any]=5						,    UpperCAmelCase_							: int=True						,    UpperCAmelCase_							: Union[str, Any]=25						,    )   ->List[Any]:



														"""simple docstring"""



														snake_case_          =						compute_mauve(
														    p_text=UpperCAmelCase_						,    q_text=UpperCAmelCase_						,    p_features=UpperCAmelCase_						,    q_features=UpperCAmelCase_						,    p_tokens=UpperCAmelCase_						,    q_tokens=UpperCAmelCase_						,    num_buckets=UpperCAmelCase_						,    pca_max_data=UpperCAmelCase_						,    kmeans_explained_var=UpperCAmelCase_						,    kmeans_num_redo=UpperCAmelCase_						,    kmeans_max_iter=UpperCAmelCase_						,    featurize_model_name=UpperCAmelCase_						,    device_id=UpperCAmelCase_						,    max_text_length=UpperCAmelCase_						,    divergence_curve_discretization_size=UpperCAmelCase_						,    mauve_scaling_factor=UpperCAmelCase_						,    verbose=UpperCAmelCase_						,    seed=UpperCAmelCase_						,    )
														return out



