support for additional training metadata (RasaHQ#5743)

* support for additional training metadata

* added test for additional training data attributes

Co-authored-by: Tom Bocklisch <[email protected]>
Co-authored-by: Roberto <[email protected]>

Author: 3 people

changelog/5743.enhancement.rst

@@ -0,0 +1,4 @@
+Support for additional training metadata.
+
+Training data messages now to support kwargs and the Rasa JSON data reader
+includes all fields when instantiating a training data instance.

rasa/nlu/training_data/formats/rasa.py

@@ -52,7 +52,7 @@ def read_from_json(self, js: Dict[Text, Any], **_) -> "TrainingData":
         all_examples = common_examples + intent_examples + entity_examples
         training_examples = []
         for ex in all_examples:
-            msg = Message.build(ex["text"], ex.get("intent"), ex.get("entities"))
+            msg = Message.build(**ex)
             training_examples.append(msg)
 
         return TrainingData(

rasa/nlu/training_data/message.py

@@ -13,11 +13,12 @@
 
 class Message:
     def __init__(
-        self, text: Text, data=None, output_properties=None, time=None
+        self, text: Text, data=None, output_properties=None, time=None, **kwargs
     ) -> None:
         self.text = text
         self.time = time
         self.data = data if data else {}
+        self.data.update(**kwargs)
 
         if output_properties:
             self.output_properties = output_properties
@@ -72,7 +73,7 @@ def __hash__(self) -> int:
         return hash((self.text, str(ordered(self.data))))
 
     @classmethod
-    def build(cls, text, intent=None, entities=None) -> "Message":
+    def build(cls, text, intent=None, entities=None, **kwargs) -> "Message":
         data = {}
         if intent:
             split_intent, response_key = cls.separate_intent_response_key(intent)
@@ -81,7 +82,7 @@ def build(cls, text, intent=None, entities=None) -> "Message":
                 data[RESPONSE_KEY_ATTRIBUTE] = response_key
         if entities:
             data[ENTITIES] = entities
-        return cls(text, data)
+        return cls(text, data, **kwargs)
 
     def get_combined_intent_response_key(self) -> Text:
         """Get intent as it appears in training data"""

tests/core/test_tracker_stores.py

@@ -343,7 +343,7 @@ def test_get_db_url_with_query():
     )
 
 
-def test_db_url_with_query_from_endpoint_config():
+def test_db_url_with_query_from_endpoint_config(tmp_path):
     endpoint_config = """
     tracker_store:
       dialect: postgresql
@@ -356,11 +356,9 @@ def test_db_url_with_query_from_endpoint_config():
         driver: my-driver
         another: query
     """
-
-    with tempfile.NamedTemporaryFile("w+", suffix="_tmp_config_file.yml") as f:
-        f.write(endpoint_config)
-        f.flush()
-        store_config = read_endpoint_config(f.name, "tracker_store")
+    f = tmp_path / "tmp_config_file.yml"
+    f.write_text(endpoint_config)
+    store_config = read_endpoint_config(str(f), "tracker_store")
 
     url = SQLTrackerStore.get_db_url(**store_config.kwargs)
 

tests/nlu/test_config.py

@@ -24,15 +24,14 @@ def test_blank_config(blank_config):
     assert final_config.as_dict() == blank_config.as_dict()
 
 
-def test_invalid_config_json():
+def test_invalid_config_json(tmp_path):
     file_config = """pipeline: [pretrained_embeddings_spacy"""  # invalid yaml
 
-    with tempfile.NamedTemporaryFile("w+", suffix="_tmp_config_file.json") as f:
-        f.write(file_config)
-        f.flush()
+    f = tmp_path / "tmp_config_file.json"
+    f.write_text(file_config)
 
-        with pytest.raises(config.InvalidConfigError):
-            config.load(f.name)
+    with pytest.raises(config.InvalidConfigError):
+        config.load(str(f))
 
 
 def test_invalid_pipeline_template():

tests/nlu/training_data/test_training_data.py

@@ -259,7 +259,7 @@ def test_markdown_single_sections():
     assert td_syn_only.entity_synonyms == {"Chines": "chinese", "Chinese": "chinese"}
 
 
-def test_repeated_entities():
+def test_repeated_entities(tmp_path):
     data = """
 {
   "rasa_nlu_data": {
@@ -279,21 +279,20 @@ def test_repeated_entities():
     ]
   }
 }"""
-    with tempfile.NamedTemporaryFile(suffix="_tmp_training_data.json") as f:
-        f.write(data.encode(io_utils.DEFAULT_ENCODING))
-        f.flush()
-        td = training_data.load_data(f.name)
-        assert len(td.entity_examples) == 1
-        example = td.entity_examples[0]
-        entities = example.get("entities")
-        assert len(entities) == 1
-        tokens = WhitespaceTokenizer().tokenize(example, attribute=TEXT)
-        start, end = MitieEntityExtractor.find_entity(entities[0], example.text, tokens)
-        assert start == 9
-        assert end == 10
-
-
-def test_multiword_entities():
+    f = tmp_path / "tmp_training_data.json"
+    f.write_text(data, io_utils.DEFAULT_ENCODING)
+    td = training_data.load_data(str(f))
+    assert len(td.entity_examples) == 1
+    example = td.entity_examples[0]
+    entities = example.get("entities")
+    assert len(entities) == 1
+    tokens = WhitespaceTokenizer().tokenize(example, attribute=TEXT)
+    start, end = MitieEntityExtractor.find_entity(entities[0], example.text, tokens)
+    assert start == 9
+    assert end == 10
+
+
+def test_multiword_entities(tmp_path):
     data = """
 {
   "rasa_nlu_data": {
@@ -313,21 +312,20 @@ def test_multiword_entities():
     ]
   }
 }"""
-    with tempfile.NamedTemporaryFile(suffix="_tmp_training_data.json") as f:
-        f.write(data.encode(io_utils.DEFAULT_ENCODING))
-        f.flush()
-        td = training_data.load_data(f.name)
-        assert len(td.entity_examples) == 1
-        example = td.entity_examples[0]
-        entities = example.get("entities")
-        assert len(entities) == 1
-        tokens = WhitespaceTokenizer().tokenize(example, attribute=TEXT)
-        start, end = MitieEntityExtractor.find_entity(entities[0], example.text, tokens)
-        assert start == 4
-        assert end == 7
-
-
-def test_nonascii_entities():
+    f = tmp_path / "tmp_training_data.json"
+    f.write_text(data, io_utils.DEFAULT_ENCODING)
+    td = training_data.load_data(str(f))
+    assert len(td.entity_examples) == 1
+    example = td.entity_examples[0]
+    entities = example.get("entities")
+    assert len(entities) == 1
+    tokens = WhitespaceTokenizer().tokenize(example, attribute=TEXT)
+    start, end = MitieEntityExtractor.find_entity(entities[0], example.text, tokens)
+    assert start == 4
+    assert end == 7
+
+
+def test_nonascii_entities(tmp_path):
     data = """
 {
   "luis_schema_version": "5.0",
@@ -345,22 +343,21 @@ def test_nonascii_entities():
     }
   ]
 }"""
-    with tempfile.NamedTemporaryFile(suffix="_tmp_training_data.json") as f:
-        f.write(data.encode(io_utils.DEFAULT_ENCODING))
-        f.flush()
-        td = training_data.load_data(f.name)
-        assert len(td.entity_examples) == 1
-        example = td.entity_examples[0]
-        entities = example.get("entities")
-        assert len(entities) == 1
-        entity = entities[0]
-        assert entity["value"] == "ßäæ ?€ö)"
-        assert entity["start"] == 19
-        assert entity["end"] == 27
-        assert entity["entity"] == "description"
-
-
-def test_entities_synonyms():
+    f = tmp_path / "tmp_training_data.json"
+    f.write_text(data, io_utils.DEFAULT_ENCODING)
+    td = training_data.load_data(str(f))
+    assert len(td.entity_examples) == 1
+    example = td.entity_examples[0]
+    entities = example.get("entities")
+    assert len(entities) == 1
+    entity = entities[0]
+    assert entity["value"] == "ßäæ ?€ö)"
+    assert entity["start"] == 19
+    assert entity["end"] == 27
+    assert entity["entity"] == "description"
+
+
+def test_entities_synonyms(tmp_path):
     data = """
 {
   "rasa_nlu_data": {
@@ -398,11 +395,10 @@ def test_entities_synonyms():
     ]
   }
 }"""
-    with tempfile.NamedTemporaryFile(suffix="_tmp_training_data.json") as f:
-        f.write(data.encode(io_utils.DEFAULT_ENCODING))
-        f.flush()
-        td = training_data.load_data(f.name)
-        assert td.entity_synonyms["New York City"] == "nyc"
+    f = tmp_path / "tmp_training_data.json"
+    f.write_text(data, io_utils.DEFAULT_ENCODING)
+    td = training_data.load_data(str(f))
+    assert td.entity_synonyms["New York City"] == "nyc"
 
 
 def cmp_message_list(firsts, seconds):
@@ -531,3 +527,24 @@ def test_load_data_from_non_existing_file():
 
 def test_is_empty():
     assert TrainingData().is_empty()
+
+
+def test_custom_attributes(tmp_path):
+    data = """
+{
+  "rasa_nlu_data": {
+    "common_examples" : [
+      {
+        "intent": "happy",
+        "text": "I'm happy.",
+        "sentiment": 0.8
+      }
+    ]
+  }
+}"""
+    f = tmp_path / "tmp_training_data.json"
+    f.write_text(data, io_utils.DEFAULT_ENCODING)
+    td = training_data.load_data(str(f))
+    assert len(td.training_examples) == 1
+    example = td.training_examples[0]
+    assert example.get("sentiment") == 0.8