Repository: OpenNMT/OpenNMT-tf
Pull request #535 (2019-10-29 17:05:49)

Make injection of special sequence tokens (
, 
) configurable
This is especially needed for language models where these special tokens may or may not be wanted depending on the task requirement.

For example, to allow nonconditional and unbounded generation of text, one would enable `
` but disable `
`. However, to generate a fixed sequence from a prefix/context, one would disable `
` but enable `
`.

This change also enables this configuration for seq2seq models (source and target). There could be reasons to want `
` at the end of the source sequences.


Patch:
diff --git a/docs/configuration.md b/docs/configuration.md
index 003cf5bb..8211d8f6 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -59,6 +59,16 @@ data:
     case_insensitive: True
     trainable: False
 
+  # (optional) For language models, configure sequence control tokens (usually
+  # represented as 
 and 
). For example, enabling "start" and disabling "end"
+  # allows nonconditional and unbounded generation (default: start=false, end=true).
+  #
+  # Advanced users could also configure this parameter for seq2seq models with e.g.
+  # source_sequence_controls and target_sequence_controls.
+  sequence_controls:
+    start: false
+    end: true
+
   # (optional) For sequence tagging tasks, the tagging scheme that is used (e.g. BIOES).
   # For supported schemes, additional evaluation metrics could be computed such as
   # precision, recall, etc. (accepted values: bioes; default: null).
diff --git a/opennmt/inputters/__init__.py b/opennmt/inputters/__init__.py
index 1d777c0e..81bd31e6 100644
--- a/opennmt/inputters/__init__.py
+++ b/opennmt/inputters/__init__.py
@@ -19,4 +19,5 @@ from opennmt.inputters.text_inputter import CharConvEmbedder
 from opennmt.inputters.text_inputter import CharRNNEmbedder
 from opennmt.inputters.text_inputter import TextInputter
 from opennmt.inputters.text_inputter import WordEmbedder
+from opennmt.inputters.text_inputter import add_sequence_controls
 from opennmt.inputters.text_inputter import load_pretrained_embeddings
diff --git a/opennmt/inputters/inputter.py b/opennmt/inputters/inputter.py
index 906f9959..f4297a0a 100644
--- a/opennmt/inputters/inputter.py
+++ b/opennmt/inputters/inputter.py
@@ -106,8 +106,18 @@ class Inputter(tf.keras.layers.Layer):
     """Returns the input signature of this inputter."""
     raise NotImplementedError()
 
-  def get_length(self, features):
-    """Returns the length of the input features, if defined."""
+  def get_length(self, features, ignore_special_tokens=False):
+    """Returns the length of the input features, if defined.
+
+    Args:
+      features: The dictionary of input features.
+      ignore_special_tokens: Ignore special tokens that were added by the
+        inputter (e.g. 
 and/or 
).
+
+    Returns:
+      The length.
+    """
+    _ = ignore_special_tokens
     return features.get("length")
 
   @abc.abstractmethod
@@ -271,14 +281,15 @@ class ParallelInputter(MultiInputter):
     else:
       return tuple(inputter.input_signature() for inputter in self.inputters)
 
-  def get_length(self, features):
+  def get_length(self, features, ignore_special_tokens=False):
     lengths = []
     for i, inputter in enumerate(self.inputters):
       if self.combine_features:
         sub_features = misc.extract_prefixed_keys(features, "inputter_{}_".format(i))
       else:
         sub_features = features[i]
-      lengths.append(inputter.get_length(sub_features))
+      lengths.append(inputter.get_length(
+          sub_features, ignore_special_tokens=ignore_special_tokens))
     if self.reducer is None:
       return lengths
     else:
@@ -378,8 +389,8 @@ class MixedInputter(MultiInputter):
       signature.update(inputter.input_signature())
     return signature
 
-  def get_length(self, features):
-    return self.inputters[0].get_length(features)
+  def get_length(self, features, ignore_special_tokens=False):
+    return self.inputters[0].get_length(features, ignore_special_tokens=ignore_special_tokens)
 
   def make_features(self, element=None, features=None, training=None):
     if features is None:
diff --git a/opennmt/inputters/text_inputter.py b/opennmt/inputters/text_inputter.py
index 1b38d6af..acae7d1b 100644
--- a/opennmt/inputters/text_inputter.py
+++ b/opennmt/inputters/text_inputter.py
@@ -160,6 +160,46 @@ def load_pretrained_embeddings(embedding_file,
 
   return pretrained
 
+def add_sequence_controls(ids, length, start_id=None, end_id=None):
+  """Adds sequence control tokens.
+
+  Args:
+    ids: Sequence of ids as 1D or 2D (batch) tensor.
+    length: Sequence length as 0D or 1D (batch) tensor.
+    start_id: Id to prepend to the sequence (set ``None`` to disable).
+    end_id: Id to append to the sequence (set ``None`` to disable).
+
+  Returns:
+    A tuple ``(ids, length)``.
+  """
+  rank = ids.shape.rank
+  if rank not in (1, 2):
+    raise ValueError("Unsupported rank %d (expected 1 or 2)" % rank)
+  batch_size = tf.shape(ids)[0] if rank == 2 else None
+
+  def _make_column(value):
+    value = tf.constant(value, dtype=ids.dtype)
+    if batch_size is not None:
+      value = tf.fill([batch_size], value)
+    return tf.expand_dims(value, -1)
+
+  if start_id is not None:
+    start_ids = _make_column(constants.START_OF_SENTENCE_ID)
+    ids = tf.concat([start_ids, ids], axis=-1)
+    length += 1
+
+  if end_id is not None:
+    end_ids = _make_column(constants.END_OF_SENTENCE_ID)
+    if batch_size is not None:
+      # Run concat on RaggedTensor to handle sequences with variable length.
+      ids = tf.RaggedTensor.from_tensor(ids, lengths=length)
+    ids = tf.concat([ids, end_ids], axis=-1)
+    if batch_size is not None:
+      ids = ids.to_tensor()
+    length += 1
+
+  return ids, length
+
 def _get_field(config, key, prefix=None, default=None, required=False):
   if prefix:
     key = "%s%s" % (prefix, key)
@@ -291,6 +331,38 @@ class WordEmbedder(TextInputter):
     self.embedding_size = embedding_size
     self.embedding_file = None
     self.dropout = dropout
+    self.decoder_mode = False
+    self.mark_start = None
+    self.mark_end = None
+
+  def set_decoder_mode(self, enable=True, mark_start=None, mark_end=None):
+    """Make this inputter produce sequences for a decoder.
+
+    In this mode, the returned "ids_out" feature is the decoder output sequence
+    and "ids" is the decoder input sequence.
+
+    Args:
+      enable: Enable the decoder mode.
+      mark_start: Mark the sequence start. If ``None``, keep the current value.
+      mark_end: Mark the sequence end. If ``None``, keep the current value.
+    """
+    self.decoder_mode = enable
+    if mark_start is not None:
+      self.mark_start = mark_start
+    if mark_end is not None:
+      self.mark_end = mark_end
+
+  def get_length(self, features, ignore_special_tokens=False):
+    length = features["length"]
+    if ignore_special_tokens:
+      # Decoder mode shifts the sequences by one timesteps.
+      num_special_tokens = -1 if self.decoder_mode else 0
+      if self.mark_start:
+        num_special_tokens += 1
+      if self.mark_end:
+        num_special_tokens += 1
+      length -= num_special_tokens
+    return length
 
   def initialize(self, data_config, asset_prefix=""):
     super(WordEmbedder, self).initialize(data_config, asset_prefix=asset_prefix)
@@ -302,14 +374,27 @@ class WordEmbedder(TextInputter):
       self.trainable = embedding.get("trainable", True)
       self.embedding_file_with_header = embedding.get("with_header", True)
       self.case_insensitive_embeddings = embedding.get("case_insensitive", True)
+    sequence_controls = _get_field(data_config, "sequence_controls", prefix=asset_prefix)
+    if sequence_controls:
+      self.mark_start = sequence_controls["start"]
+      self.mark_end = sequence_controls["end"]
 
   def make_features(self, element=None, features=None, training=None):
     """Converts words tokens to ids."""
     features = super(WordEmbedder, self).make_features(
         element=element, features=features, training=training)
-    if "ids" in features:
-      return features
-    features["ids"] = self.tokens_to_ids.lookup(features["tokens"])
+    if "ids" not in features:
+      features["ids"] = self.tokens_to_ids.lookup(features["tokens"])
+      if self.mark_start or self.mark_end:
+        features["ids"], features["length"] = add_sequence_controls(
+            features["ids"],
+            features["length"],
+            start_id=constants.START_OF_SENTENCE_ID if self.mark_start else None,
+            end_id=constants.END_OF_SENTENCE_ID if self.mark_end else None)
+    if self.decoder_mode:
+      features["ids_out"] = features["ids"][1:]
+      features["ids"] = features["ids"][:-1]
+      features["length"] -= 1
     return features
 
   def build(self, input_shape):
diff --git a/opennmt/models/language_model.py b/opennmt/models/language_model.py
index bf9df802..f2d1ebf8 100644
--- a/opennmt/models/language_model.py
+++ b/opennmt/models/language_model.py
@@ -2,7 +2,6 @@
 
 import tensorflow as tf
 
-from opennmt import constants
 from opennmt import inputters
 from opennmt import layers
 from opennmt.data import dataset as dataset_util
@@ -13,7 +12,7 @@ from opennmt.utils import misc
 
 
 class LanguageModel(model.SequenceGenerator):
-  """An experimental language model."""
+  """A language model."""
 
   def __init__(self,
                decoder,
@@ -61,16 +60,24 @@ class LanguageModel(model.SequenceGenerator):
     ids, length = features["ids"], features["length"]
     if labels is not None:
       # For training and evaluation, forward the full sequence.
-      logits, _ = self._decode(ids, length, training=training)
+      logits, _ = self._decode(
+          labels.get("ids", ids),
+          labels.get("length", length),
+          training=training)
       outputs = dict(logits=logits)
     else:
       assert_fixed_length = tf.debugging.Assert(
           tf.reduce_all(tf.equal(length, tf.reduce_max(length))),
           ["Language model does not support variable length contexts during "
            "generation, consider setting batch_size or length_bucket_width to 1"])
-
-      # Run decoder one the context, if any.
-      with tf.control_dependencies([assert_fixed_length]):
+      assert_non_empty_start = tf.debugging.Assert(
+          tf.math.not_equal(tf.math.reduce_max(length), 0),
+          ["The language model requires a context sequence to initialize the decoding. "
+           "If you want nonconditional sequence generation, you should configure the "
+           "sequence_controls parameter before training."])
+
+      # Run decoder on the context, if any.
+      with tf.control_dependencies([assert_fixed_length, assert_non_empty_start]):
         context_ids, start_ids = tf.split(ids, [tf.shape(ids)[1] - 1, 1], axis=1)
         context_length = length - 1
         batch_size = tf.shape(context_length)[0]
@@ -96,6 +103,10 @@ class LanguageModel(model.SequenceGenerator):
       sampled_length = tf.reshape(sampled_length, [batch_size])
 
       # Build the full prediction.
+      if self.features_inputter.mark_start:
+        # Remove leading 
 if included in the context sequence.
+        ids = ids[:, 1:]
+        length -= 1
       full_ids = tf.concat([ids, sampled_ids], 1)
       full_length = length + sampled_length
       tokens = self.features_inputter.ids_to_tokens.lookup(full_ids)
@@ -133,15 +144,29 @@ class LanguageModelInputter(inputters.WordEmbedder):
   input sequence.
   """
 
+  def initialize(self, data_config, asset_prefix=""):
+    super(LanguageModelInputter, self).initialize(data_config, asset_prefix=asset_prefix)
+    # Set default sequence controls for backward compatibility.
+    if self.mark_start is None:
+      self.mark_start = False
+    if self.mark_end is None:
+      self.mark_end = True
+
   def make_features(self, element=None, features=None, training=None):
+    base_features = features if features is not None else {}
+
+    # Features define the decoder context during inference. As the context is a prefix,
+    # we should disable the end sequence control token.
+    saved_mark_end = self.mark_end
+    self.set_decoder_mode(enable=False, mark_end=False)
     features = super(LanguageModelInputter, self).make_features(
-        element=element, features=features, training=training)
-    labels = {
-        "ids_out": tf.concat([features["ids"][1:], [constants.END_OF_SENTENCE_ID]], 0),
-        "length": tf.identity(features["length"])
-    }
-    if not training:
-      labels["tokens"] = tf.concat([features["tokens"][1:], [constants.END_OF_SENTENCE_TOKEN]], 0)
+        element=element, features=base_features.copy(), training=training)
+
+    # Labels define the decoder input/output sequences during training and evaluation.
+    self.set_decoder_mode(enable=True, mark_end=saved_mark_end)
+    labels = super(LanguageModelInputter, self).make_features(
+        element=element, features=base_features.copy(), training=training)
+
     return features, labels
 
   def make_inference_dataset(self,
diff --git a/opennmt/models/model.py b/opennmt/models/model.py
index 2dc58e21..81418abc 100644
--- a/opennmt/models/model.py
+++ b/opennmt/models/model.py
@@ -376,6 +376,13 @@ class Model(tf.keras.layers.Layer):
 class SequenceGenerator(Model):
   """Base class for models generating sequences."""
 
+  @property
+  def decoder_inputter(self):
+    """The inputter used on the decoder side."""
+    return (
+        self.labels_inputter if not self.unsupervised
+        else self.features_inputter)
+
   def score(self, features, labels):
     outputs, _ = self(features, labels=labels)
     cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
@@ -387,7 +394,7 @@ class SequenceGenerator(Model):
         "cross_entropy": cross_entropy,
         "score": scores,
         "tokens": labels["tokens"],
-        "length": labels["length"] - 1  # -1 for the special token.
+        "length": self.decoder_inputter.get_length(labels, ignore_special_tokens=True)
     }
     if "attention" in outputs:
       results["attention"] = outputs["attention"]
@@ -397,11 +404,8 @@ class SequenceGenerator(Model):
     if params is None:
       params = {}
     length = score["length"]
-    output_tokenizer = (
-        self.labels_inputter.tokenizer if not self.unsupervised
-        else self.features_inputter.tokenizer)
     tokens = score["tokens"][:length]
-    sentence = output_tokenizer.detokenize(tokens)
+    sentence = self.decoder_inputter.tokenizer.detokenize(tokens)
     token_level_scores = None
     attention = None
     if params.get("with_token_level"):
diff --git a/opennmt/models/sequence_to_sequence.py b/opennmt/models/sequence_to_sequence.py
index c309ca89..5ca7679c 100644
--- a/opennmt/models/sequence_to_sequence.py
+++ b/opennmt/models/sequence_to_sequence.py
@@ -314,11 +314,10 @@ class SequenceToSequence(model.SequenceGenerator):
           labels["noisy_ids_out"],
           labels["noisy_length"],
           eta=params.get("max_margin_eta", 0.1))
-    labels_lengths = self.labels_inputter.get_length(labels)
     loss, loss_normalizer, loss_token_normalizer = losses.cross_entropy_sequence_loss(
         logits,
         labels["ids_out"],
-        labels_lengths,
+        labels["length"],
         label_smoothing=params.get("label_smoothing", 0.0),
         average_in_time=params.get("average_loss_in_time", False),
         training=training)
@@ -333,7 +332,7 @@ class SequenceToSequence(model.SequenceGenerator):
           loss += losses.guided_alignment_cost(
               attention[:, :-1],  # Do not constrain last timestep.
               gold_alignments,
-              sequence_length=labels_lengths - 1,
+              sequence_length=self.labels_inputter.get_length(labels, ignore_special_tokens=True),
               cost_type=guided_alignment_type,
               weight=params.get("guided_alignment_weight", 1))
     return loss, loss_normalizer, loss_token_normalizer
@@ -417,6 +416,7 @@ class SequenceToSequenceInputter(inputters.ExampleInputter):
                share_parameters=False):
     super(SequenceToSequenceInputter, self).__init__(
         features_inputter, labels_inputter, share_parameters=share_parameters)
+    labels_inputter.set_decoder_mode(mark_start=True, mark_end=True)
     self.alignment_file = None
 
   def initialize(self, data_config, asset_prefix=""):
@@ -440,22 +440,11 @@ class SequenceToSequenceInputter(inputters.ExampleInputter):
     if alignment is not None:
       labels["alignment"] = text.alignment_matrix_from_pharaoh(
           alignment,
-          self.features_inputter.get_length(features),
-          self.labels_inputter.get_length(labels))
-    _shift_target_sequence(labels)
-    if "noisy_ids" in labels:
-      _shift_target_sequence(labels, prefix="noisy_")
+          self.features_inputter.get_length(features, ignore_special_tokens=True),
+          self.labels_inputter.get_length(labels, ignore_special_tokens=True))
     return features, labels
 
 
-def _shift_target_sequence(labels, prefix=""):
-  labels_ids = labels["%sids" % prefix]
-  bos = tf.constant([constants.START_OF_SENTENCE_ID], dtype=labels_ids.dtype)
-  eos = tf.constant([constants.END_OF_SENTENCE_ID], dtype=labels_ids.dtype)
-  labels["%sids" % prefix] = tf.concat([bos, labels_ids], axis=0)
-  labels["%sids_out" % prefix] = tf.concat([labels_ids, eos], axis=0)
-  labels["%slength" % prefix] += 1
-
 def align_tokens_from_attention(tokens, attention):
   """Returns aligned tokens from the attention.
 


Test patch:
diff --git a/opennmt/tests/inputter_test.py b/opennmt/tests/inputter_test.py
index 5654db1e..294b5660 100644
--- a/opennmt/tests/inputter_test.py
+++ b/opennmt/tests/inputter_test.py
@@ -19,6 +19,7 @@ from opennmt.data import dataset as dataset_util
 from opennmt.data import noise
 from opennmt.inputters import inputter, text_inputter, record_inputter
 from opennmt.layers import reducer
+from opennmt.tests import test_util
 from opennmt.utils.misc import item_or_tuple, count_lines
 
 
@@ -131,6 +132,18 @@ class InputterTest(tf.test.TestCase):
     self.assertAllEqual([1, 1], embeddings[1])
     self.assertAllEqual([3, 3], embeddings[2])
 
+  @parameterized.expand([
+      [[3, 4], 2, 1, 2, [1, 3, 4, 2], 4],
+      [[[3, 4], [5, 6]], [2, 1], 1, None, [[1, 3, 4], [1, 5, 6]], [3, 2]],
+      [[[3, 4], [5, 6]], [2, 1], None, 2, [[3, 4, 2], [5, 2, 0]], [3, 2]],
+  ])
+  def testAddSequenceControls(self, ids, length, start_id, end_id, expected_ids, expected_length):
+    ids = tf.constant(ids, dtype=tf.int64)
+    length = tf.constant(length, dtype=tf.int32)
+    ids, length = inputters.add_sequence_controls(ids, length, start_id=start_id, end_id=end_id)
+    self.assertAllEqual(self.evaluate(ids), expected_ids)
+    self.assertAllEqual(self.evaluate(length), expected_length)
+
   def _checkFeatures(self, features, expected_shapes):
     for name, expected_shape in six.iteritems(expected_shapes):
       self.assertIn(name, features)
@@ -173,6 +186,18 @@ class InputterTest(tf.test.TestCase):
     self.assertAllEqual([[2, 1, 4]], features["ids"])
     self.assertAllEqual([1, 3, 10], transformed.shape)
 
+  def testWordEmbedderForDecoder(self):
+    vocab_file = test_util.make_vocab(
+        os.path.join(self.get_temp_dir(), "vocab.txt"), ["the", "world", "hello", "toto"])
+    embedder = text_inputter.WordEmbedder(embedding_size=10)
+    embedder.set_decoder_mode(mark_start=True, mark_end=True)
+    embedder.initialize({"vocabulary": vocab_file})
+    features = self.evaluate(embedder.make_features(tf.constant("hello world")))
+    self.assertEqual(features["length"], 3)
+    self.assertEqual(embedder.get_length(features, ignore_special_tokens=True), 2)
+    self.assertAllEqual(features["ids"], [1, 5, 4])
+    self.assertAllEqual(features["ids_out"], [5, 4, 2])
+
   def testWordEmbedderWithTokenizer(self):
     vocab_file = self._makeTextFile("vocab.txt", ["the", "world", "hello", "￭"])
     data_file = self._makeTextFile("data.txt", ["hello world!"])
diff --git a/opennmt/tests/model_test.py b/opennmt/tests/model_test.py
index 4f2067c9..8ccc4fe4 100644
--- a/opennmt/tests/model_test.py
+++ b/opennmt/tests/model_test.py
@@ -288,6 +288,55 @@ class ModelTest(tf.test.TestCase):
         prediction_heads=["tokens", "length"],
         params=params)
 
+  def testLanguageModelInputter(self):
+    vocabulary_path = test_util.make_vocab(
+        os.path.join(self.get_temp_dir(), "vocab.txt"), ["a", "b", "c"])
+
+    inputter = models.LanguageModelInputter(embedding_size=10)
+    inputter.initialize({
+        "vocabulary": vocabulary_path,
+        "sequence_controls": {"start": True, "end": False}})
+    features, labels = self.evaluate(inputter.make_features(tf.constant("a b c")))
+    self.assertAllEqual(features["ids"], [1, 3, 4, 5])
+    self.assertEqual(features["length"], 4)
+    self.assertAllEqual(labels["ids"], [1, 3, 4])
+    self.assertAllEqual(labels["ids_out"], [3, 4, 5])
+    self.assertEqual(labels["length"], 3)
+
+    # Backward compatibility mode.
+    inputter = models.LanguageModelInputter(embedding_size=10)
+    inputter.initialize({"vocabulary": vocabulary_path})
+    features, labels = self.evaluate(inputter.make_features(tf.constant("a b c")))
+    self.assertAllEqual(features["ids"], [3, 4, 5])
+    self.assertEqual(features["length"], 3)
+    self.assertAllEqual(labels["ids"], [3, 4, 5])
+    self.assertAllEqual(labels["ids_out"], [4, 5, 2])
+    self.assertEqual(labels["length"], 3)
+
+  def testLanguageModelWithMissingStart(self):
+    _, data_config = self._makeToyLMData()
+    decoder = decoders.SelfAttentionDecoder(
+        2, num_units=16, num_heads=4, ffn_inner_dim=32, num_sources=0)
+    model = models.LanguageModel(decoder, embedding_size=16)
+    model.initialize(data_config)
+    features, _ = model.features_inputter.make_features(tf.constant(""))
+    with self.assertRaises(tf.errors.InvalidArgumentError):
+      model(features)
+
+  def testLanguageModelWithStartOfSentence(self):
+    _, data_config = self._makeToyLMData()
+    data_config["sequence_controls"] = dict(start=True, end=False)
+    decoder = decoders.SelfAttentionDecoder(
+        2, num_units=16, num_heads=4, ffn_inner_dim=32, num_sources=0)
+    model = models.LanguageModel(decoder, embedding_size=16)
+    model.initialize(data_config, params={"maximum_decoding_length": 1})
+    features, _ = model.features_inputter.make_features(tf.constant(""))
+    features = tf.nest.map_structure(lambda t: tf.expand_dims(t, 0), features)  # Add batch dim.
+    _, predictions = self.evaluate(model(features))
+    # Predictions should not include the leading 
.
+    self.assertEqual(predictions["length"][0], 1)
+    self.assertTupleEqual(predictions["tokens"].shape, (1, 1))
+
   @parameterized.expand([
       [tf.estimator.ModeKeys.TRAIN],
       [tf.estimator.ModeKeys.EVAL],
diff --git a/opennmt/tests/test_util.py b/opennmt/tests/test_util.py
index 0fcc6400..4093b672 100644
--- a/opennmt/tests/test_util.py
+++ b/opennmt/tests/test_util.py
@@ -17,6 +17,16 @@ def make_data_file(path, lines):
       data.write("%s\n" % line)
   return path
 
+def make_vocab(path, tokens):
+  vocabulary = vocab.Vocab(special_tokens=[
+      constants.PADDING_TOKEN,
+      constants.START_OF_SENTENCE_TOKEN,
+      constants.END_OF_SENTENCE_TOKEN])
+  for token in tokens:
+    vocabulary.add(token)
+  vocabulary.serialize(path)
+  return path
+
 def make_vocab_from_file(path, data_file):
   vocabulary = vocab.Vocab(special_tokens=[
       constants.PADDING_TOKEN,