def test_process(
text: Text,
lookup: List[Dict[Text, List[Text]]],
expected_entities: List[Dict[Text, Any]],
):
message = Message(data={TEXT: text})
training_data = TrainingData()
training_data.lookup_tables = lookup
training_data.training_examples = [
Message(data={
TEXT: "Hi Max!",
"entities": [{
"entity": "person",
"value": "Max"
}]
}),
Message(
data={
TEXT: "I live in Berlin",
"entities": [{
"entity": "city",
"value": "Berlin"
}],
}),
]
entity_extractor = RegexEntityExtractor()
entity_extractor.train(training_data)
entity_extractor.process(message)
entities = message.get(ENTITIES)
assert entities == expected_entities
def test_lowercase(
text: Text,
case_sensitive: bool,
lookup: List[Dict[Text, List[Text]]],
expected_entities: List[Dict[Text, Any]],
):
message = Message(data={TEXT: text})
training_data = TrainingData()
training_data.lookup_tables = lookup
training_data.training_examples = [
Message(
data={
TEXT: "Hi Max!",
INTENT: "greet",
ENTITIES: [{"entity": "person", "value": "Max"}],
}
),
Message(
data={
TEXT: "I live in Berlin",
INTENT: "inform",
ENTITIES: [{"entity": "city", "value": "Berlin"}],
}
),
]
entity_extractor = RegexEntityExtractor({"case_sensitive": case_sensitive})
entity_extractor.train(training_data)
entity_extractor.process(message)
entities = message.get(ENTITIES)
assert entities == expected_entities
def test_do_not_overwrite_any_entities():
message = Message(data={TEXT: "Max lives in Berlin.", INTENT: "infrom"})
message.set(ENTITIES, [{
"entity": "person",
"value": "Max",
"start": 0,
"end": 3
}])
training_data = TrainingData()
training_data.training_examples = [
Message(
data={
TEXT: "Hi Max!",
INTENT: "greet",
ENTITIES: [{
"entity": "person",
"value": "Max"
}],
}),
Message(
data={
TEXT: "I live in Berlin",
INTENT: "inform",
ENTITIES: [{
"entity": "city",
"value": "Berlin"
}],
}),
]
training_data.lookup_tables = [{
"name":
"city",
"elements": ["London", "Berlin", "Amsterdam"]
}]
entity_extractor = RegexEntityExtractor()
entity_extractor.train(training_data)
entity_extractor.process(message)
entities = message.get(ENTITIES)
assert entities == [
{
"entity": "person",
"value": "Max",
"start": 0,
"end": 3
},
{
"entity": "city",
"value": "Berlin",
"start": 13,
"end": 19,
"extractor": "RegexEntityExtractor",
},
]