Merge pull request #5565 from RasaHQ/mask-vs-sequence

Mask vs sequence

rasa/nlu/classifiers/diet_classifier.py

@@ -87,10 +87,10 @@
 
 TEXT_FEATURES = f"{TEXT}_features"
 LABEL_FEATURES = f"{LABEL}_features"
-TEXT_MASK = f"{TEXT}_mask"
-LABEL_MASK = f"{LABEL}_mask"
 LABEL_IDS = f"{LABEL}_ids"
 TAG_IDS = "tag_ids"
+TEXT_SEQ_LENGTH = f"{TEXT}_lengths"
+LABEL_SEQ_LENGTH = f"{LABEL}_lengths"
 
 
 class DIETClassifier(IntentClassifier, EntityExtractor):
@@ -484,7 +484,7 @@ def _create_label_data(
         # to track correctly dynamic sequences
         label_data.add_features(LABEL_IDS, [np.expand_dims(label_ids, -1)])
 
-        label_data.add_mask(LABEL_MASK, LABEL_FEATURES)
+        label_data.add_lengths(LABEL_SEQ_LENGTH, LABEL_FEATURES)
 
         return label_data
 
@@ -558,8 +558,8 @@ def _create_model_data(
         model_data.add_features(LABEL_IDS, [np.expand_dims(label_ids, -1)])
         model_data.add_features(TAG_IDS, [tag_ids])
 
-        model_data.add_mask(TEXT_MASK, TEXT_FEATURES)
-        model_data.add_mask(LABEL_MASK, LABEL_FEATURES)
+        model_data.add_lengths(TEXT_SEQ_LENGTH, TEXT_FEATURES)
+        model_data.add_lengths(LABEL_SEQ_LENGTH, LABEL_FEATURES)
 
         return model_data
 
@@ -1165,10 +1165,6 @@ def _prepare_entity_recognition_layers(self) -> None:
             average="micro",
         )
 
-    @staticmethod
-    def _get_sequence_lengths(mask: tf.Tensor) -> tf.Tensor:
-        return tf.cast(tf.reduce_sum(mask[:, :, 0], axis=1), tf.int32)
-
     def _combine_sparse_dense_features(
         self,
         features: List[Union[np.ndarray, tf.Tensor, tf.SparseTensor]],
@@ -1258,10 +1254,14 @@ def _create_sequence(
 
     def _create_all_labels(self) -> Tuple[tf.Tensor, tf.Tensor]:
         all_label_ids = self.tf_label_data[LABEL_IDS][0]
+
+        label_lengths = self.sequence_lengths_for(
+            self.tf_label_data[LABEL_SEQ_LENGTH][0]
+        )
+        mask_label = self._compute_mask(label_lengths)
+
         x = self._create_bow(
-            self.tf_label_data[LABEL_FEATURES],
-            self.tf_label_data[LABEL_MASK][0],
-            self.label_name,
+            self.tf_label_data[LABEL_FEATURES], mask_label, self.label_name,
         )
         all_labels_embed = self._tf_layers[f"embed.{LABEL}"](x)
 
@@ -1355,13 +1355,23 @@ def _calculate_entity_loss(
 
         return loss, f1
 
+    @staticmethod
+    def _compute_mask(sequence_lengths: tf.Tensor) -> tf.Tensor:
+        mask = tf.sequence_mask(sequence_lengths, dtype=tf.float32)
+        # explicitly add last dimension to mask
+        # to track correctly dynamic sequences
+        return tf.expand_dims(mask, -1)
+
+    def sequence_lengths_for(self, sequence_lengths: tf.Tensor) -> tf.Tensor:
+        return tf.cast(sequence_lengths, dtype=tf.int32)
+
     def batch_loss(
         self, batch_in: Union[Tuple[tf.Tensor], Tuple[np.ndarray]]
     ) -> tf.Tensor:
         tf_batch_data = self.batch_to_model_data_format(batch_in, self.data_signature)
 
-        mask_text = tf_batch_data[TEXT_MASK][0]
-        sequence_lengths = self._get_sequence_lengths(mask_text)
+        sequence_lengths = self.sequence_lengths_for(tf_batch_data[TEXT_SEQ_LENGTH][0])
+        mask_text = self._compute_mask(sequence_lengths)
 
         (
             text_transformed,
@@ -1390,11 +1400,14 @@ def batch_loss(
             # get _cls_ vector for intent classification
             cls = self._last_token(text_transformed, sequence_lengths)
 
+            label_lengths = self.sequence_lengths_for(
+                tf_batch_data[LABEL_SEQ_LENGTH][0]
+            )
+            mask_label = self._compute_mask(label_lengths)
+
             label_ids = tf_batch_data[LABEL_IDS][0]
             label = self._create_bow(
-                tf_batch_data[LABEL_FEATURES],
-                tf_batch_data[LABEL_MASK][0],
-                self.label_name,
+                tf_batch_data[LABEL_FEATURES], mask_label, self.label_name,
             )
             loss, acc = self._calculate_label_loss(cls, label, label_ids)
             self.intent_loss.update_state(loss)
@@ -1420,8 +1433,8 @@ def batch_predict(
             batch_in, self.predict_data_signature
         )
 
-        mask_text = tf_batch_data[TEXT_MASK][0]
-        sequence_lengths = self._get_sequence_lengths(mask_text)
+        sequence_lengths = self.sequence_lengths_for(tf_batch_data[TEXT_SEQ_LENGTH][0])
+        mask_text = self._compute_mask(sequence_lengths)
 
         text_transformed, _, _, _ = self._create_sequence(
             tf_batch_data[TEXT_FEATURES], mask_text, self.text_name

rasa/nlu/selectors/response_selector.py

@@ -17,9 +17,9 @@
     DIET,
     TEXT_FEATURES,
     LABEL_FEATURES,
-    TEXT_MASK,
-    LABEL_MASK,
     LABEL_IDS,
+    TEXT_SEQ_LENGTH,
+    LABEL_SEQ_LENGTH,
 )
 from rasa.utils.tensorflow.constants import (
     LABEL,
@@ -432,8 +432,10 @@ def _prepare_layers(self) -> None:
     def _create_all_labels(self) -> Tuple[tf.Tensor, tf.Tensor]:
         all_label_ids = self.tf_label_data[LABEL_IDS][0]
 
-        mask_label = self.tf_label_data[LABEL_MASK][0]
-        sequence_lengths_label = self._get_sequence_lengths(mask_label)
+        sequence_lengths_label = self.sequence_lengths_for(
+            self.tf_label_data[LABEL_SEQ_LENGTH][0]
+        )
+        mask_label = self._compute_mask(sequence_lengths_label)
 
         label_transformed, _, _, _ = self._create_sequence(
             self.tf_label_data[LABEL_FEATURES], mask_label, self.label_name
@@ -449,8 +451,10 @@ def batch_loss(
     ) -> tf.Tensor:
         tf_batch_data = self.batch_to_model_data_format(batch_in, self.data_signature)
 
-        mask_text = tf_batch_data[TEXT_MASK][0]
-        sequence_lengths_text = self._get_sequence_lengths(mask_text)
+        sequence_lengths_text = self.sequence_lengths_for(
+            tf_batch_data[TEXT_SEQ_LENGTH][0]
+        )
+        mask_text = self._compute_mask(sequence_lengths_text)
 
         (
             text_transformed,
@@ -465,8 +469,10 @@ def batch_loss(
             sequence_ids=True,
         )
 
-        mask_label = tf_batch_data[LABEL_MASK][0]
-        sequence_lengths_label = self._get_sequence_lengths(mask_label)
+        sequence_lengths_label = self.sequence_lengths_for(
+            tf_batch_data[LABEL_SEQ_LENGTH][0]
+        )
+        mask_label = self._compute_mask(sequence_lengths_label)
 
         label_transformed, _, _, _ = self._create_sequence(
             tf_batch_data[LABEL_FEATURES], mask_label, self.label_name
@@ -506,8 +512,10 @@ def batch_predict(
             batch_in, self.predict_data_signature
         )
 
-        mask_text = tf_batch_data[TEXT_MASK][0]
-        sequence_lengths_text = self._get_sequence_lengths(mask_text)
+        sequence_lengths_text = self.sequence_lengths_for(
+            tf_batch_data[TEXT_SEQ_LENGTH][0]
+        )
+        mask_text = self._compute_mask(sequence_lengths_text)
 
         text_transformed, _, _, _ = self._create_sequence(
             tf_batch_data[TEXT_FEATURES], mask_text, self.text_name

rasa/utils/tensorflow/model_data.py

@@ -144,20 +144,17 @@ def add_features(self, key: Text, features: List[np.ndarray]):
         # update number of examples
         self.num_examples = self.number_of_examples()
 
-    def add_mask(self, key: Text, from_key: Text):
-        """Calculate mask for given key and put it under specified key."""
-
+    def add_lengths(self, key: Text, from_key: Text) -> None:
+        """Adds np.array of lengths of sequences to data under given key."""
         if not self.data.get(from_key):
             return
 
         self.data[key] = []
 
         for data in self.data[from_key]:
             if data.size > 0:
-                # explicitly add last dimension to mask
-                # to track correctly dynamic sequences
-                mask = np.array([np.ones((x.shape[0], 1)) for x in data])
-                self.data[key].append(mask)
+                lengths = np.array([x.shape[0] for x in data])
+                self.data[key].append(lengths)
                 break
 
     def split(