transformer_fast_decode, beam search: take an optional cache and return it

Some models, eg. semantic parsing models with copying mechanisms,
want to use the output of Transformer for multiple predictions.
One way to do so is to modify the symbols_to_logits_fn to generate
the additional predictions and save it in the cache dictionary.

To do so, though, fast_decode() must allow an externally supplied cache,
and must return it to the caller after the loop.

Author: Giovanni Campagna

tensor2tensor/layers/latent_layers.py

@@ -134,7 +134,7 @@ def symbols_to_logits_fn(ids):
 
   initial_ids = tf.zeros([tf.shape(latents_dense_in)[0]], dtype=tf.int32)
   length = tf.shape(latents_dense_in)[1]
-  ids, _ = beam_search.beam_search(
+  ids, _, _ = beam_search.beam_search(
       symbols_to_logits_fn,
       initial_ids,
       1,

tensor2tensor/models/research/transformer_nat.py

@@ -227,7 +227,7 @@ def symbols_to_logits_fn(ids):
 
   initial_ids = tf.zeros([tf.shape(latents_dense_in)[0]], dtype=tf.int32)
   length = tf.shape(latents_dense_in)[1]
-  ids, _ = beam_search.beam_search(
+  ids, _, _ = beam_search.beam_search(
       symbols_to_logits_fn,
       initial_ids,
       beam_size=1,

tensor2tensor/models/research/transformer_vae.py

@@ -286,7 +286,7 @@ def symbols_to_logits_fn(ids):
 
   initial_ids = tf.zeros([tf.shape(latents_dense_in)[0]], dtype=tf.int32)
   length = tf.shape(latents_dense_in)[1]
-  ids, _ = beam_search.beam_search(
+  ids, _, _ = beam_search.beam_search(
       symbols_to_logits_fn, initial_ids, beam_size, length,
       vocab_size, alpha=0.0, eos_id=-1, stop_early=False)
 

tensor2tensor/models/transformer.py

@@ -812,7 +812,8 @@ def fast_decode(encoder_output,
                 eos_id=beam_search.EOS_ID,
                 batch_size=None,
                 force_decode_length=False,
-                scope_prefix="body/"):
+                scope_prefix="body/",
+                cache=None):
   """Given encoder output and a symbols to logits function, does fast decoding.
 
   Implements both greedy and beam search decoding, uses beam search iff
@@ -859,7 +860,9 @@ def fast_decode(encoder_output,
   vars_3d_num_heads = (
       hparams.num_heads if hparams.get("attention_variables_3d") else 0)
 
-  cache = {
+  if cache is None:
+    cache = dict()
+  cache.update({
       "layer_%d" % layer: {
           "k":
               common_attention.split_heads(
@@ -870,7 +873,7 @@ def fast_decode(encoder_output,
           "f":
               tf.zeros([batch_size, 0, hparams.hidden_size]),
       } for layer in range(num_layers)
-  }
+  })
 
   if encoder_output is not None:
     for layer in range(num_layers):
@@ -894,7 +897,7 @@ def fast_decode(encoder_output,
 
   if beam_size > 1:  # Beam Search
     initial_ids = sos_id * tf.ones([batch_size], dtype=tf.int32)
-    decoded_ids, scores = beam_search.beam_search(
+    decoded_ids, scores, cache = beam_search.beam_search(
         symbols_to_logits_fn,
         initial_ids,
         beam_size,
@@ -940,7 +943,7 @@ def is_not_finished(i, hit_eos, *_):
     hit_eos = tf.fill([batch_size], False)
     next_id = sos_id * tf.ones([batch_size, 1], dtype=tf.int64)
     initial_log_prob = tf.zeros([batch_size], dtype=tf.float32)
-    _, _, _, decoded_ids, _, log_prob = tf.while_loop(
+    _, _, _, decoded_ids, cache, log_prob = tf.while_loop(
         is_not_finished,
         inner_loop, [
             tf.constant(0), hit_eos, next_id, decoded_ids, cache,
@@ -956,7 +959,7 @@ def is_not_finished(i, hit_eos, *_):
         ])
     scores = log_prob
 
-  return {"outputs": decoded_ids, "scores": scores}
+  return {"outputs": decoded_ids, "scores": scores, "cache": cache}
 
 
 @registry.register_model

tensor2tensor/utils/beam_search.py

@@ -505,7 +505,7 @@ def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
         tf.less(i, decode_length), tf.logical_not(bound_is_met))
 
   (_, alive_seq, alive_log_probs, finished_seq, finished_scores,
-   finished_flags, _) = tf.while_loop(
+   finished_flags, states) = tf.while_loop(
        _is_finished,
        inner_loop, [
            tf.constant(0), alive_seq, alive_log_probs, finished_seq,
@@ -535,4 +535,4 @@ def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
       tf.reduce_any(finished_flags, 1), finished_seq, alive_seq)
   finished_scores = tf.where(
       tf.reduce_any(finished_flags, 1), finished_scores, alive_log_probs)
-  return finished_seq, finished_scores
+  return finished_seq, finished_scores, states

tensor2tensor/utils/beam_search_test.py

@@ -37,7 +37,7 @@ def symbols_to_logits(_):
       # Just return random logits
       return tf.random_uniform((batch_size * beam_size, vocab_size))
 
-    final_ids, final_probs = beam_search.beam_search(
+    final_ids, final_probs, _ = beam_search.beam_search(
         symbols_to_logits, initial_ids, beam_size, decode_length, vocab_size,
         0.)
 
@@ -113,7 +113,7 @@ def symbols_to_logits(ids):
       logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
       return logits
 
-    final_ids, final_probs = beam_search.beam_search(
+    final_ids, final_probs, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
@@ -144,7 +144,7 @@ def symbols_to_logits(ids):
       logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
       return logits
 
-    final_ids, final_probs = beam_search.beam_search(
+    final_ids, final_probs, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
@@ -173,7 +173,7 @@ def symbols_to_logits(ids):
       logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
       return logits
 
-    final_ids, final_probs = beam_search.beam_search(
+    final_ids, final_probs, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
@@ -213,7 +213,7 @@ def symbols_to_logits(ids):
       logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
       return logits
 
-    final_ids, final_scores = beam_search.beam_search(
+    final_ids, final_scores, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
@@ -256,7 +256,7 @@ def symbols_to_logits(ids):
       return logits
 
     # Disable early stopping
-    final_ids, final_scores = beam_search.beam_search(
+    final_ids, final_scores, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
@@ -302,7 +302,7 @@ def symbols_to_logits(ids, _, states):
     states["state"] = tf.placeholder_with_default(
         states["state"], shape=(None, 1))
 
-    final_ids, _ = beam_search.beam_search(
+    final_ids, _, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,
@@ -319,6 +319,41 @@ def symbols_to_logits(ids, _, states):
       except tf.errors.InvalidArgumentError as e:
         raise AssertionError(e.message)
 
+  def testStatesAfterLoop(self):
+    batch_size = 1
+    beam_size = 1
+    vocab_size = 2
+    decode_length = 3
+
+    initial_ids = tf.constant([0] * batch_size)  # GO
+    probabilities = tf.constant([[[0.7, 0.3]], [[0.4, 0.6]], [[0.5, 0.5]]])
+
+    def symbols_to_logits(ids, _, states):
+      pos = tf.shape(ids)[1] - 1
+      logits = tf.to_float(tf.log(probabilities[pos, :]))
+      states["state"] += 1
+      return logits, states
+
+    states = {
+        "state": tf.zeros((batch_size, 1)),
+    }
+    states["state"] = tf.placeholder_with_default(
+        states["state"], shape=(None, 1))
+
+    _, _, final_states = beam_search.beam_search(
+        symbols_to_logits,
+        initial_ids,
+        beam_size,
+        decode_length,
+        vocab_size,
+        0.0,
+        eos_id=1,
+        states=states)
+    
+    with self.test_session() as sess:
+      final_states = sess.run(final_states)
+    self.assertAllEqual([[1]], final_states["state"])
+
   def testStateBeamTwo(self):
     batch_size = 1
     beam_size = 2
@@ -352,7 +387,7 @@ def symbols_to_logits(ids, _, states):
     states["state"] = tf.placeholder_with_default(
         states["state"], shape=(None, 1))
 
-    final_ids, _ = beam_search.beam_search(
+    final_ids, _, _ = beam_search.beam_search(
         symbols_to_logits,
         initial_ids,
         beam_size,

tensor2tensor/utils/t2t_model.py

@@ -698,7 +698,7 @@ def symbols_to_logits_fn(ids):
       inputs = features["inputs"]
       decode_length = (common_layers.shape_list(inputs)[1] +
                        features.get("decode_length", decode_length))
-    ids, scores = beam_search.beam_search(
+    ids, scores, _ = beam_search.beam_search(
         symbols_to_logits_fn,
         initial_ids,
         beam_size,