import comet  # From: unbabel-comet
import torch

import datasets


A						=			datasets.logging.get_logger(__name__)

A						=			'\\n@inproceedings{rei-EtAl:2020:WMT,\n   author    = {Rei, Ricardo  and  Stewart, Craig  and  Farinha, Ana C  and  Lavie, Alon},\n   title     = {Unbabel\'s Participation in the WMT20 Metrics Shared Task},\n   booktitle      = {Proceedings of the Fifth Conference on Machine Translation},\n   month          = {November},\n   year           = {2020},\n   address        = {Online},\n   publisher      = {Association for Computational Linguistics},\n   pages     = {909--918},\n}\n@inproceedings{rei-etal-2020-comet,\n   title = "{COMET}: A Neural Framework for {MT} Evaluation",\n   author = "Rei, Ricardo  and\n      Stewart, Craig  and\n      Farinha, Ana C  and\n      Lavie, Alon",\n   booktitle = "Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)",\n   month = nov,\n   year = "2020",\n   address = "Online",\n   publisher = "Association for Computational Linguistics",\n   url = "https://www.aclweb.org/anthology/2020.emnlp-main.213",\n   pages = "2685--2702",\n}\n'

A						=			'\\nCrosslingual Optimized Metric for Evaluation of Translation (COMET) is an open-source framework used to train Machine Translation metrics that achieve high levels of correlation with different types of human judgments (HTER, DA\'s or MQM).\nWith the release of the framework the authors also released fully trained models that were used to compete in the WMT20 Metrics Shared Task achieving SOTA in that years competition.\n\nSee the [README.md] file at https://unbabel.github.io/COMET/html/models.html for more information.\n'

A						=			'\nCOMET score.\n\nArgs:\n\n`sources` (list of str): Source sentences\n`predictions` (list of str): candidate translations\n`references` (list of str): reference translations\n`cuda` (bool): If set to True, runs COMET using GPU\n`show_progress` (bool): Shows progress\n`model`: COMET model to be used. Will default to `wmt-large-da-estimator-1719` if None.\n\nReturns:\n    `samples`: List of dictionaries with `src`, `mt`, `ref` and `score`.\n    `scores`: List of scores.\n\nExamples:\n\n    >>> comet_metric = datasets.load_metric(\'comet\')\n    >>> # comet_metric = load_metric(\'comet\', \'wmt20-comet-da\')  # you can also choose which model to use\n    >>> source = ["Dem Feuer konnte Einhalt geboten werden", "Schulen und Kindergärten wurden eröffnet."]\n    >>> hypothesis = ["The fire could be stopped", "Schools and kindergartens were open"]\n    >>> reference = ["They were able to control the fire.", "Schools and kindergartens opened"]\n    >>> results = comet_metric.compute(predictions=hypothesis, references=reference, sources=source)\n    >>> print([round(v, 2) for v in results["scores"]])\n    [0.19, 0.92]\n'





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


		'''simple docstring'''






		def   _snake_case       (       self     :     Optional[int]		)      ->	List[Any]:
				return datasets.MetricInfo(
				    description=_DESCRIPTION     ,  citation=_CITATION     ,  homepage='https://unbabel.github.io/COMET/html/index.html'     ,  inputs_description=_KWARGS_DESCRIPTION     ,  features=datasets.Features(
				        {
				            'sources': datasets.Value('string'     ,  id='sequence'		),
				            'predictions': datasets.Value('string'     ,  id='sequence'		),
				            'references': datasets.Value('string'     ,  id='sequence'		),
				        }		)     ,  codebase_urls=['https://github.com/Unbabel/COMET']     ,  reference_urls=[
				        'https://github.com/Unbabel/COMET',
				        'https://www.aclweb.org/anthology/2020.emnlp-main.213/',
				        'http://www.statmt.org/wmt20/pdf/2020.wmt-1.101.pdf6',
				    ]     ,  )
		def   _snake_case       (       self     :     str     ,  snake_case__     :     str		)      ->	Optional[int]:
				if self.config_name == "default":
						_lowerCamelCase				   =  comet.load_from_checkpoint(comet.download_model('wmt20-comet-da'		)		)
				else:
						_lowerCamelCase				   =  comet.load_from_checkpoint(comet.download_model(self.config_name		)		)



		def   _snake_case       (       self     :     Dict     ,  snake_case__     :     Optional[Any]     ,  snake_case__     :     List[str]     ,  snake_case__     :     Any     ,  snake_case__     :     List[str]=None     ,  snake_case__     :     str=False		)      ->	Optional[Any]:
				if gpus is None:
						_lowerCamelCase				   =  1 if torch.cuda.is_available() else 0
				_lowerCamelCase				   =  {'src': sources, 'mt': predictions, 'ref': references}
				_lowerCamelCase				   =  [dict(zip(snake_case__     ,  snake_case__		)		) for t in zip(*data.values()		)]
				_lowerCamelCase   ,							_lowerCamelCase				   =  self.scorer.predict(snake_case__     ,  gpus=snake_case__     ,  progress_bar=snake_case__		)
				return {"mean_score": mean_score, "scores": scores}
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 	lowerCamelCase					(	UpperCamelCase     :				int							)   ->    Union[str, Any]:
		_lowerCamelCase				   =  torch.exp(UpperCamelCase							)
		_lowerCamelCase				   =  torch.sum(UpperCamelCase			,       dim=1							)  # sum of exp(x_i)
		_lowerCamelCase				   =  torch.sum(x * exp_x			,       dim=1							)  # sum of x_i * exp(x_i)
		return torch.log(UpperCamelCase							) - B / A





class     lowerCAmelCase__     (				nn.Module			):


		'''simple docstring'''






		def __init__(       self     :     Union[str, Any]     ,  snake_case__     :     Any		)      ->	Optional[Any]:
				super().__init__()
				_lowerCamelCase				   =  config.output_attentions
				_lowerCamelCase				   =  config.output_hidden_states
				_lowerCamelCase				   =  nn.ModuleList([BertLayer(snake_case__		) for _ in range(config.num_hidden_layers		)]		)
				_lowerCamelCase				   =  nn.ModuleList([BertHighway(snake_case__		) for _ in range(config.num_hidden_layers		)]		)

				_lowerCamelCase				   =  [-1 for _ in range(config.num_hidden_layers		)]
		def   _snake_case       (       self     :     Union[str, Any]     ,  snake_case__     :     str		)      ->	Optional[int]:
				if (type(snake_case__		) is float) or (type(snake_case__		) is int):
						for i in range(len(self.early_exit_entropy		)		):
								_lowerCamelCase				   =  x
				else:
						_lowerCamelCase				   =  x
		def   _snake_case       (       self     :     Tuple     ,  snake_case__     :     Optional[int]		)      ->	int:
				_lowerCamelCase				   =  pooler.state_dict()
				for highway in self.highway:
						for name, param in highway.pooler.state_dict().items():
								param.copy_(loaded_model[name]		)



		def   _snake_case       (       self     :     Union[str, Any]     ,  snake_case__     :     Union[str, Any]     ,  snake_case__     :     int=None     ,  snake_case__     :     Optional[int]=None     ,  snake_case__     :     List[str]=None     ,  snake_case__     :     Dict=None     ,  )      ->	List[str]:
				_lowerCamelCase				   =  ()
				_lowerCamelCase				   =  ()
				_lowerCamelCase				   =  ()
				for i, layer_module in enumerate(self.layer		):
						if self.output_hidden_states:
								_lowerCamelCase				   =  all_hidden_states + (hidden_states,)

						_lowerCamelCase				   =  layer_module(
						    snake_case__     ,  snake_case__     ,  head_mask[i]     ,  snake_case__     ,  snake_case__		)
						_lowerCamelCase				   =  layer_outputs[0]

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

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

						_lowerCamelCase				   =  self.highway[i](snake_case__		)
						# logits, pooled_output

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

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

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

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

				_lowerCamelCase				   =  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). '     ,SCREAMING_SNAKE_CASE_     ,)
class     lowerCAmelCase__     (				SCREAMING_SNAKE_CASE_			):


		'''simple docstring'''






		def __init__(       self     :     Dict     ,  snake_case__     :     List[Any]		)      ->	Any:
				super().__init__(snake_case__		)
				_lowerCamelCase				   =  config

				_lowerCamelCase				   =  BertEmbeddings(snake_case__		)
				_lowerCamelCase				   =  DeeBertEncoder(snake_case__		)
				_lowerCamelCase				   =  BertPooler(snake_case__		)

				self.init_weights()
		def   _snake_case       (       self     :     Tuple		)      ->	Tuple:
				self.encoder.init_highway_pooler(self.pooler		)
		def   _snake_case       (       self     :     int		)      ->	List[Any]:
				return self.embeddings.word_embeddings
		def   _snake_case       (       self     :     Optional[int]     ,  snake_case__     :     List[Any]		)      ->	Any:
				_lowerCamelCase				   =  value
		def   _snake_case       (       self     :     List[Any]     ,  snake_case__     :     int		)      ->	Any:
				for layer, heads in heads_to_prune.items():
						self.encoder.layer[layer].attention.prune_heads(snake_case__		)



		@add_start_docstrings_to_model_forward(snake_case__		)
		def   _snake_case       (       self     :     Union[str, Any]     ,  snake_case__     :     Dict=None     ,  snake_case__     :     Any=None     ,  snake_case__     :     Dict=None     ,  snake_case__     :     Any=None     ,  snake_case__     :     List[str]=None     ,  snake_case__     :     Union[str, Any]=None     ,  snake_case__     :     str=None     ,  snake_case__     :     List[Any]=None     ,  )      ->	List[str]:
				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:
						_lowerCamelCase				   =  input_ids.size()
				elif inputs_embeds is not None:
						_lowerCamelCase				   =  inputs_embeds.size()[:-1]
				else:
						raise ValueError('You have to specify either input_ids or inputs_embeds'		)

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

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

				# 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.
				_lowerCamelCase				   =  self.get_extended_attention_mask(snake_case__     ,  snake_case__     ,  snake_case__		)

				# 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:
						_lowerCamelCase				   =  encoder_attention_mask[:, None, :, :]
				if encoder_attention_mask.dim() == 2:
						_lowerCamelCase				   =  encoder_attention_mask[:, None, None, :]

				_lowerCamelCase				   =  encoder_extended_attention_mask.to(
				    dtype=next(self.parameters()		).dtype		)  # fp16 compatibility
				_lowerCamelCase				   =  (1.0 - encoder_extended_attention_mask) * -1_0000.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]
				_lowerCamelCase				   =  self.get_head_mask(snake_case__     ,  self.config.num_hidden_layers		)

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

				_lowerCamelCase				   =  (
				    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     lowerCAmelCase__     (				SCREAMING_SNAKE_CASE_			):


		'''simple docstring'''






		def __init__(       self     :     Optional[int]     ,  snake_case__     :     Dict     ,  snake_case__     :     Any		)      ->	str:
				_lowerCamelCase				   =  message
				_lowerCamelCase				   =  exit_layer  # start from 1!





class     lowerCAmelCase__     (				nn.Module			):


		'''simple docstring'''






		def __init__(       self     :     Union[str, Any]     ,  snake_case__     :     Optional[Any]		)      ->	Any:
				super().__init__()
				_lowerCamelCase				   =  BertPooler(snake_case__		)
				_lowerCamelCase				   =  nn.Dropout(config.hidden_dropout_prob		)
				_lowerCamelCase				   =  nn.Linear(config.hidden_size     ,  config.num_labels		)



		def   _snake_case       (       self     :     Optional[Any]     ,  snake_case__     :     Dict		)      ->	Optional[Any]:
				# Pooler
				_lowerCamelCase				   =  encoder_outputs[0]
				_lowerCamelCase				   =  self.pooler(snake_case__		)
				# "return" pooler_output

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

				# Dropout and classification
				_lowerCamelCase				   =  bmodel_output[1]

				_lowerCamelCase				   =  self.dropout(snake_case__		)
				_lowerCamelCase				   =  self.classifier(snake_case__		)

				return logits, pooled_output





@add_start_docstrings(
    'Bert Model (with early exiting - DeeBERT) with a classifier on top,\n    also takes care of multi-layer training. '     ,SCREAMING_SNAKE_CASE_     ,)
class     lowerCAmelCase__     (				SCREAMING_SNAKE_CASE_			):


		'''simple docstring'''






		def __init__(       self     :     Optional[int]     ,  snake_case__     :     Dict		)      ->	Dict:
				super().__init__(snake_case__		)
				_lowerCamelCase				   =  config.num_labels
				_lowerCamelCase				   =  config.num_hidden_layers

				_lowerCamelCase				   =  DeeBertModel(snake_case__		)
				_lowerCamelCase				   =  nn.Dropout(config.hidden_dropout_prob		)
				_lowerCamelCase				   =  nn.Linear(config.hidden_size     ,  self.config.num_labels		)

				self.init_weights()



		@add_start_docstrings_to_model_forward(snake_case__		)
		def   _snake_case       (       self     :     Union[str, Any]     ,  snake_case__     :     Union[str, Any]=None     ,  snake_case__     :     Optional[int]=None     ,  snake_case__     :     str=None     ,  snake_case__     :     Dict=None     ,  snake_case__     :     str=None     ,  snake_case__     :     Union[str, Any]=None     ,  snake_case__     :     int=None     ,  snake_case__     :     Union[str, Any]=-1     ,  snake_case__     :     Any=False     ,  )      ->	Union[str, Any]:
				_lowerCamelCase				   =  self.num_layers
				try:
						_lowerCamelCase				   =  self.bert(
						    snake_case__     ,  attention_mask=snake_case__     ,  token_type_ids=snake_case__     ,  position_ids=snake_case__     ,  head_mask=snake_case__     ,  inputs_embeds=snake_case__     ,  )
						# sequence_output, pooled_output, (hidden_states), (attentions), highway exits

						_lowerCamelCase				   =  outputs[1]

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

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

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

						if train_highway:
								_lowerCamelCase				   =  (sum(highway_losses[:-1]		),) + outputs
								# exclude the final highway, of course
						else:
								_lowerCamelCase				   =  (loss,) + outputs
				if not self.training:
						_lowerCamelCase				   =  outputs + ((original_entropy, highway_entropy), exit_layer)
						if output_layer >= 0:
								_lowerCamelCase				   =  (
								    (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)