Merge pull request #348 from medicode/set_precision_recall

Set precision and recall metrics

Author: lukaszkaiser

tensor2tensor/utils/metrics.py

@@ -43,7 +43,8 @@ class Metrics(object):
   ROUGE_2_F = "rouge_2_fscore"
   ROUGE_L_F = "rouge_L_fscore"
   EDIT_DISTANCE = "edit_distance"
-
+  SET_PRECISION = 'set_precision'
+  SET_RECALL = 'set_recall'
 
 def padded_rmse(predictions, labels, weights_fn=common_layers.weights_all):
   predictions, labels = common_layers.pad_with_zeros(predictions, labels)
@@ -188,6 +189,48 @@ def padded_accuracy(predictions,
     padded_labels = tf.to_int32(padded_labels)
     return tf.to_float(tf.equal(outputs, padded_labels)), weights
 
+def set_precision(predictions,
+                  labels,
+                  weights_fn=common_layers.weights_nonzero):
+  """Precision of set predictions.
+
+  Args:
+    predictions : A Tensor of scores of shape (batch, nlabels)
+    labels: A Tensor of int32s giving true set elements of shape (batch, seq_length)
+
+  Returns:
+    hits: A Tensor of shape (batch, nlabels)
+    weights: A Tensor of shape (batch, nlabels)
+  """
+  with tf.variable_scope("set_precision", values=[predictions, labels]):
+    labels = tf.squeeze(labels, [2, 3])
+    labels = tf.one_hot(labels, predictions.shape[-1])
+    labels = tf.reduce_max(labels, axis=1)
+    labels = tf.cast(labels, tf.bool)
+    predictions = predictions > 0
+    return tf.to_float(tf.equal(labels, predictions)), tf.to_float(predictions)
+  
+def set_recall(predictions,
+                    labels,
+                    weights_fn=common_layers.weights_nonzero):
+  """Recall of set predictions.
+
+  Args:
+    predictions : A Tensor of scores of shape (batch, nlabels)
+    labels: A Tensor of int32s giving true set elements of shape (batch, seq_length)
+
+  Returns:
+    hits: A Tensor of shape (batch, nlabels)
+    weights: A Tensor of shape (batch, nlabels)
+  """
+  with tf.variable_scope("set_recall", values=[predictions, labels]):
+    labels = tf.squeeze(labels, [2, 3])
+    labels = tf.one_hot(labels, predictions.shape[-1])
+    labels = tf.reduce_max(labels, axis=1)
+    labels = tf.cast(labels, tf.bool)
+    predictions = predictions > 0
+    return tf.to_float(tf.equal(labels, predictions)), tf.to_float(labels)
+
 
 def create_evaluation_metrics(problems, model_hparams):
   """Creates the evaluation metrics for the model.
@@ -278,4 +321,6 @@ def wrapped_metric_fn():
     Metrics.ROUGE_2_F: rouge.rouge_2_fscore,
     Metrics.ROUGE_L_F: rouge.rouge_l_fscore,
     Metrics.EDIT_DISTANCE: sequence_edit_distance,
+    Metrics.SET_PRECISION: set_precision,
+    Metrics.SET_RECALL: set_recall,
 }