text: from typing import Optional, List
  
  from overrides import overrides
  import torch
  
  from stog.metrics.metric import Metric
  
  
  class AttachmentScores(Metric):
      """
      Computes labeled and unlabeled attachment scores for a
      dependency parse, as well as sentence level exact match
      for both labeled and unlabeled trees. Note that the input
      to this metric is the sampled predictions, not the distribution
      itself.
      Parameters
      ----------
      ignore_classes : ``List[int]``, optional (default = None)
          A list of label ids to ignore when computing metrics.
      """
      def __init__(self, ignore_classes: List[int] = None) -> None:
          self._labeled_correct = 0.
          self._unlabeled_correct = 0.
          self._exact_labeled_correct = 0.
          self._exact_unlabeled_correct = 0.
          self._total_words = 0.
          self._total_sentences = 0.
  
          self._total_loss = 0.
          self._total_edge_node_loss = 0.
          self._total_edge_label_loss = 0.
  
          self._ignore_classes: List[int] = ignore_classes or []
  
      def __call__(self, # type: ignore
                   predicted_indices: torch.Tensor,
                   predicted_labels: torch.Tensor,
                   gold_indices: torch.Tensor,
                   gold_labels: torch.Tensor,
                   mask: Optional[torch.Tensor] = None,
                   edge_node_loss = 0,
                   edge_label_loss = 0):
          """
          Parameters
          ----------
          predicted_indices : ``torch.Tensor``, required.
              A tensor of head index predictions of shape (batch_size, timesteps).
          predicted_labels : ``torch.Tensor``, required.
              A tensor of arc label predictions of shape (batch_size, timesteps).
          gold_indices : ``torch.Tensor``, required.
              A tensor of the same shape as ``predicted_indices``.
          gold_labels : ``torch.Tensor``, required.
              A tensor of the same shape as ``predicted_labels``.
          mask: ``torch.Tensor``, optional (default = None).
              A tensor of the same shape as ``predicted_indices``.
          """
          unwrapped = self.unwrap_to_tensors(predicted_indices, predicted_labels,
                                             gold_indices, gold_labels, mask)
          predicted_indices, predicted_labels, gold_indices, gold_labels, mask = unwrapped
  
          mask = mask.long()
          predicted_indices = predicted_indices.long()
          predicted_labels = predicted_labels.long()
          gold_indices = gold_indices.long()
          gold_labels = gold_labels.long()
  
          # Multiply by a mask donoting locations of
          # gold labels which we should ignore.
          for label in self._ignore_classes:
              label_mask = gold_labels.eq(label)
              mask = mask * (1 - label_mask).long()
  
          correct_indices = predicted_indices.eq(gold_indices).long() * mask
          unlabeled_exact_match = (correct_indices + (1 - mask)).prod(dim=-1)
          correct_labels = predicted_labels.eq(gold_labels).long() * mask
          correct_labels_and_indices = correct_indices * correct_labels
          labeled_exact_match = (correct_labels_and_indices + (1 - mask)).prod(dim=-1)
  
          self._unlabeled_correct += correct_indices.sum()
          self._exact_unlabeled_correct += unlabeled_exact_match.sum()
          self._labeled_correct += correct_labels_and_indices.sum()
          self._exact_labeled_correct += labeled_exact_match.sum()
          self._total_sentences += correct_indices.size(0)
          self._total_words += correct_indices.numel() - (1 - mask).sum()
  
          self._total_loss += edge_node_loss + edge_label_loss
          self._total_edge_node_loss += edge_node_loss
          self._total_edge_label_loss += edge_label_loss
  
      def get_metric(self, reset: bool = False):
          """
          Returns
          -------
          The accumulated metrics as a dictionary.
          """
          unlabeled_attachment_score = 0.0
          labeled_attachment_score = 0.0
          unlabeled_exact_match = 0.0
          labeled_exact_match = 0.0
          edge_loss = 0.0
          edge_node_loss = 0.0
          edge_label_loss = 0.0
          if self._total_words > 0.0:
              unlabeled_attachment_score = float(self._unlabeled_correct) / float(self._total_words)
              labeled_attachment_score = float(self._labeled_correct) / float(self._total_words)
              edge_loss = float(self._total_loss) / float(self._total_words)
              edge_node_loss = float(self._total_edge_node_loss) / float(self._total_words)
              edge_label_loss = float(self._total_edge_label_loss) / float(self._total_words)
          if self._total_sentences > 0:
              unlabeled_exact_match = float(self._exact_unlabeled_correct) / float(self._total_sentences)
              labeled_exact_match = float(self._exact_labeled_correct) / float(self._total_sentences)
          if reset:
              self.reset()
          return {
              "UAS": unlabeled_attachment_score,
              "LAS": labeled_attachment_score,
              # "UEM": unlabeled_exact_match,
              # "LEM": labeled_exact_match,
              "EL": edge_loss
          }
  
  
      def reset(self):
          self._labeled_correct = 0.
          self._unlabeled_correct = 0.
          self._exact_labeled_correct = 0.
          self._exact_unlabeled_correct = 0.
          self._total_words = 0.
          self._total_sentences = 0.
  
          self._total_loss = 0.
          self._total_edge_node_loss = 0.
          self._total_edge_label_loss = 0.
  