@Model.register('jain_wallace_attention_binary_classifier')
class JWAED(Model):
  '''
    Encoder/decoder with attention model for binary classification as described in 'Attention is Not
    Explanation' (Jain and Wallace 2019) - Jain and Wallace) 2019 (https://arxiv.org/pdf/1902.10186.pdf)
  '''
  def __init__(self, vocab: Vocabulary, word_embeddings: TextFieldEmbedder, encoder: Seq2SeqEncoder, 
               decoder: FeedForward):
    super().__init__(vocab)
    self.word_embeddings = word_embeddings
    self.num_classes = self.vocab.get_vocab_size('labels')
    self.encoder = encoder
    self.attention = AdditiveAttention(encoder.get_output_dim())
    self.decoder = decoder
    self.metrics = {
      'accuracy': CategoricalAccuracy(),
      'f1_measure': F1Measure(positive_label=1),
      'auc': Auc(positive_label=1)
    }
    self.loss = torch.nn.BCEWithLogitsLoss()

  @overrides
  def forward(self, tokens: Dict[str, torch.LongTensor], label: torch.LongTensor = None) -> Dict[str, torch.Tensor]:

    output_dict = {}

    # encode
    tokens_mask = util.get_text_field_mask(tokens)
    embedded_tokens = self.word_embeddings(tokens)
    encoded_tokens = self.encoder(embedded_tokens, tokens_mask)

    # compute attention
    attention = self.attention(encoded_tokens, tokens_mask)
    context = (attention.unsqueeze(-1) * encoded_tokens).sum(1)

    # decode
    logits = self.decoder(context)
    positive_probability = torch.sigmoid(logits)
    class_probabilities = torch.cat((1 - positive_probability, positive_probability), dim=1)
    output_dict['class_probabilities'] = class_probabilities

    if label is not None:
      loss = self.loss(logits, label.unsqueeze(-1).float())
      output_dict['loss'] = loss
      self.metrics['accuracy'](class_probabilities, label)
      self.metrics['f1_measure'](class_probabilities, label)
      self.metrics['auc'](positive_probability.squeeze(-1), label)

    return output_dict

  @overrides
  def get_metrics(self, reset: bool = False) -> Dict[str, float]:
    precision, recall, f1_measure = self.metrics['f1_measure'].get_metric(reset=reset)
    return {
      # f1 get_metric returns (precision, recall, f1)
      'positive/precision': precision,
      'positive/recall': recall,
      'f1_measure': f1_measure,
      'accuracy': self.metrics['accuracy'].get_metric(reset=reset),
      'auc': self.metrics['auc'].get_metric(reset=reset)
    }

  @overrides
  def decode(self, output_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    '''
    Does a simple argmax over the class probabilities, converts indices to string labels, and
    adds a ``'label'`` key to the dictionary with the result.
    '''
    class_probabilities = output_dict['class_probabilities']
    predictions = class_probabilities.cpu().data.numpy()
    argmax_indices = np.argmax(predictions, axis=-1)
    labels = [self.vocab.get_token_from_index(x, namespace='labels') for x in argmax_indices]
    output_dict['label'] = labels
    return output_dict

    archive = load_archive(model_path)
    predictor = Predictor.from_archive(archive, 'text_classifier')
    validation_instances = predictor._dataset_reader.read('/.../datasets/SST/test.txt')
    gradients, outputs = predictor.get_gradients(validation_instances)
    print(gradients)
    # prints {}