internal merge of PR tensorflow#1242

PiperOrigin-RevId: 223252032

Author: stefan-falk

tensor2tensor/data_generators/speech_recognition.py

@@ -140,6 +140,6 @@ def preprocess_example(self, example, mode, hparams):
   def eval_metrics(self):
     defaults = super(SpeechRecognitionProblem, self).eval_metrics()
     return defaults + [
-      metrics.Metrics.EDIT_DISTANCE,
-      metrics.Metrics.WORD_ERROR_RATE
+        metrics.Metrics.EDIT_DISTANCE,
+        metrics.Metrics.WORD_ERROR_RATE
     ]

tensor2tensor/utils/metrics.py

@@ -681,38 +681,38 @@ def metric_means():
   return metric_accum, metric_means
 
 
-def word_error_rate(raw_predictions, labels, lookup=None,
+def word_error_rate(raw_predictions,
+                    labels,
+                    lookup=None,
                     weights_fn=common_layers.weights_nonzero):
-  """
-  :param raw_predictions:
-  :param labels:
-  :param lookup:
-    A tf.constant mapping indices to output tokens.
-  :param weights_fn:
-  :return:
+  """Calculate word error rate.
+
+  Args:
+    raw_predictions: The raw predictions.
+    labels: The actual labels.
+    lookup: A tf.constant mapping indices to output tokens.
+    weights_fn: Weighting function.
+
+  Returns:
     The word error rate.
   """
 
   def from_tokens(raw, lookup_):
     gathered = tf.gather(lookup_, tf.cast(raw, tf.int32))
-    joined = tf.regex_replace(tf.reduce_join(gathered, axis=1), b'<EOS>.*', b'')
-    cleaned = tf.regex_replace(joined, b'_', b' ')
-    tokens = tf.string_split(cleaned, ' ')
+    joined = tf.regex_replace(tf.reduce_join(gathered, axis=1), b"<EOS>.*", b"")
+    cleaned = tf.regex_replace(joined, b"_", b" ")
+    tokens = tf.string_split(cleaned, " ")
     return tokens
 
   def from_characters(raw, lookup_):
-    """
-    Convert ascii+2 encoded codes to string-tokens.
-    """
+    """Convert ascii+2 encoded codes to string-tokens."""
     corrected = tf.bitcast(
-      tf.clip_by_value(
-        tf.subtract(raw, 2), 0, 255
-      ), tf.uint8)
+        tf.clip_by_value(tf.subtract(raw, 2), 0, 255), tf.uint8)
 
     gathered = tf.gather(lookup_, tf.cast(corrected, tf.int32))[:, :, 0]
     joined = tf.reduce_join(gathered, axis=1)
-    cleaned = tf.regex_replace(joined, b'\0', b'')
-    tokens = tf.string_split(cleaned, ' ')
+    cleaned = tf.regex_replace(joined, b"\0", b"")
+    tokens = tf.string_split(cleaned, " ")
     return tokens
 
   if lookup is None:
@@ -727,18 +727,16 @@ def from_characters(raw, lookup_):
   with tf.variable_scope("word_error_rate", values=[raw_predictions, labels]):
 
     raw_predictions = tf.squeeze(
-      tf.argmax(raw_predictions, axis=-1), axis=(2, 3))
+        tf.argmax(raw_predictions, axis=-1), axis=(2, 3))
     labels = tf.squeeze(labels, axis=(2, 3))
 
     reference = convert_fn(labels, lookup)
     predictions = convert_fn(raw_predictions, lookup)
 
     distance = tf.reduce_sum(
-      tf.edit_distance(predictions, reference, normalize=False)
-    )
+        tf.edit_distance(predictions, reference, normalize=False))
     reference_length = tf.cast(
-      tf.size(reference.values, out_type=tf.int32), dtype=tf.float32
-    )
+        tf.size(reference.values, out_type=tf.int32), dtype=tf.float32)
 
     return distance / reference_length, reference_length
 

tensor2tensor/utils/metrics_test.py

@@ -108,18 +108,18 @@ def testSequenceEditDistanceMetric(self):
   def testWordErrorRateMetric(self):
 
     ref = np.asarray([
-      # a b c
-      [97, 34, 98, 34, 99],
-      [97, 34, 98, 34, 99],
-      [97, 34, 98, 34, 99],
-      [97, 34, 98, 34, 99],
+        # a b c
+        [97, 34, 98, 34, 99],
+        [97, 34, 98, 34, 99],
+        [97, 34, 98, 34, 99],
+        [97, 34, 98, 34, 99],
     ])
 
     hyp = np.asarray([
-      [97, 34, 98, 34, 99],  # a b c
-      [97, 34, 98, 0, 0],  # a b
-      [97, 34, 98, 34, 100],  # a b d
-      [0, 0, 0, 0, 0]  # empty
+        [97, 34, 98, 34, 99],  # a b c
+        [97, 34, 98, 0, 0],  # a b
+        [97, 34, 98, 34, 100],  # a b d
+        [0, 0, 0, 0, 0]  # empty
     ])
 
     labels = np.reshape(ref, ref.shape + (1, 1))
@@ -130,9 +130,8 @@ def testWordErrorRateMetric(self):
         predictions[i, j, 0, 0, idx] = 1
 
     with self.test_session() as session:
-      actual_wer, actual_ref_len = session.run(
-        metrics.word_error_rate(predictions, labels)
-      )
+      actual_wer, unused_actual_ref_len = session.run(
+          metrics.word_error_rate(predictions, labels))
 
     expected_wer = 0.417
     places = 3