def _extract_entities(self, message: Message) -> List[Dict[Text, Any]]:
"""Extract entities of the given type from the given user message."""
entities = []
flags = 0 # default flag
if not self.case_sensitive:
flags = re.IGNORECASE
for pattern in self.patterns:
matches = re.finditer(pattern["pattern"],
message.get(TEXT),
flags=flags)
matches = list(matches)
for match in matches:
start_index = match.start()
end_index = match.end()
entities.append({
ENTITY_ATTRIBUTE_TYPE:
pattern["name"],
ENTITY_ATTRIBUTE_START:
start_index,
ENTITY_ATTRIBUTE_END:
end_index,
ENTITY_ATTRIBUTE_VALUE:
message.get(TEXT)[start_index:end_index],
})
return entities
def test_entity_synonyms_substitute_two_entity():
example = Message(
text="Looking for a chines restaurant in New York tomorrow",
data={
"entities": [{
"entity": "type",
"value": "chinese",
"start": 14,
"end": 20
}, {
"entity": "city",
"value": "New York",
"start": 35,
"end": 43
}]
})
ent_synonyms = {"chines": "chinese", "new york": "NYC"}
EntitySynonymBegin(synonyms=ent_synonyms).process(example)
assert example.text == "Looking for a chinese restaurant in NYC tomorrow"
e_type = list(
filter(lambda e: e["entity"] == 'type', example.get("entities")))[0]
e_city = list(
filter(lambda e: e["entity"] == 'city', example.get("entities")))[0]
assert e_type["start"] == 14
assert e_type["end"] == 21
assert e_city["start"] == 36
assert e_city["end"] == 39
def process(self, message: Message, **kwargs: Any) -> None:
if not self.patterns:
return
extracted_entities = self._extract_entities(message)
extracted_entities = self.add_extractor_name(extracted_entities)
message.set(ENTITIES,
message.get(ENTITIES, []) + extracted_entities,
add_to_output=True)
def parse_training_example(example: Text, intent: Optional[Text] = None) -> "Message":
"""Extract entities and synonyms, and convert to plain text."""
entities = find_entities_in_training_example(example)
plain_text = replace_entities(example)
return Message.build(plain_text, intent, entities)
def _collect_messages(evts: List[Dict[Text, Any]]) -> List[Message]:
"""Collect the message text and parsed data from the UserMessage events
into a list"""
from rasa.nlu.extractors.duckling_http_extractor import \
DucklingHTTPExtractor
from rasa.nlu.extractors.mitie_entity_extractor import MitieEntityExtractor
from rasa.nlu.extractors.spacy_entity_extractor import SpacyEntityExtractor
msgs = []
for evt in evts:
if evt.get("event") == UserUttered.type_name:
data = evt.get("parse_data")
for entity in data.get("entities", []):
excluded_extractors = [
DucklingHTTPExtractor.__name__,
SpacyEntityExtractor.__name__,
MitieEntityExtractor.__name__
]
logger.debug("Exclude entity marking of following extractors"
" {} when writing nlu data "
"to file.".format(excluded_extractors))
if entity.get("extractor") in excluded_extractors:
data["entities"].remove(entity)
msg = Message.build(data["text"], data["intent"]["name"],
data["entities"])
msgs.append(msg)
return msgs
def test_entity_synonyms_substitute_and_replace():
initial_text = "Looking for a chines restaurant in New York tomorrow for three people"
initial_entities = [{
"entity": "type",
"value": "chines",
"start": 14,
"end": 20
}, {
"entity": "city",
"value": "New York",
"start": 35,
"end": 43
}, {
"entity": "count",
"value": "three",
"start": 57,
"end": 62
}]
example = Message(text=initial_text, data={
"entities": initial_entities,
})
ent_synonyms = {"chines": "chinese", "new york": "NYC", "three": "3"}
EntitySynonymBegin(synonyms=ent_synonyms).process(example)
EntitySynonymEnd().process(example)
def _get_example(config=None, gazette=None, primary=None):
if primary is None:
primary = {
"entity": "type",
"value": "chines",
"start": 14,
"end": 20,
"extractor": "ner_crf",
}
return _process_example(
Message(
text="Looking for a chines restaurant in New York",
data={
"entities": [
primary,
{
"entity": "type",
"value": "restaurant",
"start": 21,
"end": 31,
"extractor": "ner_crf",
},
{
"entity": "city",
"value": "New York",
"start": 35,
"end": 43,
"extractor": "ner_crf",
},
]
},
),
config=config,
gazette=gazette,
)
def test_entity_sweeper():
entities = [{
"entity": "cuisine",
"value": "chinese",
"start": 0,
"end": 6
}, {
"entity": "time",
"value": "whatever",
"start": 0,
"end": 6
}]
sweeper = Sweeper(component_config={'entity_names': ['time']})
message = Message("xxx", {'entities': entities})
sweeper.process(message)
assert len(message.get('entities')) == 1
assert message.get('entities')[0]["entity"] == "cuisine"
def test_entity_synonyms_substitute_three_entity():
example = Message(
text=
"Looking for a chines restaurant in New York tomorrow for three people",
data={
"entities": [
{
"entity": "type",
"value": "chines",
"start": 14,
"end": 20
},
{
"entity": "city",
"value": "New York",
"start": 35,
"end": 43
},
{
"entity": "count",
"value": "three",
"start": 57,
"end": 62
},
]
},
)
ent_synonyms = {"chines": "chinese", "new york": "NYC", "three": "3"}
EntitySynonymBegin(synonyms=ent_synonyms).process(example)
assert (example.text ==
"Looking for a chinese restaurant in NYC tomorrow for 3 people")
e_type = list(
filter(lambda e: e["entity"] == "type", example.get("entities")))[0]
e_city = list(
filter(lambda e: e["entity"] == "city", example.get("entities")))[0]
e_count = list(
filter(lambda e: e["entity"] == "count", example.get("entities")))[0]
assert e_type["start"] == 14
assert e_type["end"] == 21
assert e_city["start"] == 36
assert e_city["end"] == 39
assert e_count["start"] == 53
assert e_count["end"] == 54
def test_entity_synonyms_substitute_one_entity():
example = Message(text="Looking for a chines restaurant",
data={
"entities": [{
"entity": "type",
"value": "chinese",
"start": 14,
"end": 20
}]
})
ent_synonyms = {"chines": "chinese"}
EntitySynonymBegin(synonyms=ent_synonyms).process(example)
assert example.text == "Looking for a chinese restaurant"
e_type = list(
filter(lambda e: e["entity"] == 'type', example.get("entities")))[0]
assert e_type["start"] == 14
assert e_type["end"] == 21
def test_entity_synonyms_substitute():
example = Message(text="Looking for a chines restaurant in New York",
data={
"entities": [{
"entity": "type",
"value": "chinese",
"start": 14,
"end": 20
}, {
"entity": "city",
"value": "New York",
"start": 35,
"end": 43
}]
})
ent_synonyms = {"chines": "chinese", "new york": "NYC"}
EntitySynonymBegin(synonyms=ent_synonyms).process(example)
assert example.text == "Looking for a chinese restaurant in NYC"
def _setup_example(config=None):
instance = _get_instance(config=config)
message = Message(text='This is a tst message')
flagged_tokens = [
{
"offset": 10,
"token": "tst",
"type": "UnknownToken",
"suggestions": [
{
"suggestion": "test",
"score": 0.95155325585711
},
{
"suggestion": "text",
"score": 0.805342621979041
}
]
}
]
return instance, message, flagged_tokens
def test_classification(self, trained_classifier, message, intent):
text = Message(message)
trained_classifier.process(text)
assert text.get("intent").get("name", "NOT_CLASSIFIED") == intent
def test_entity_synonyms_substitute_and_replace_w_insertions():
text_initial = "Looking for a chines restaurant in New York tomorrow for three people"
initial_entities = [{
"entity": "type",
"value": "chines",
"start": 14,
"end": 20
}, {
"entity": "city",
"value": "New York",
"start": 35,
"end": 43
}, {
"entity": "count",
"value": "three",
"start": 57,
"end": 62
}]
example = Message(text=text_initial, data={
"entities": initial_entities,
})
ent_synonyms = {"chines": "chinese", "new york": "NYC", "three": "3"}
EntitySynonymBegin(synonyms=ent_synonyms).process(example)
# import IPython
# IPython.embed()
example.data["entities"].extend([
{
"entity": "action",
"value": "Looking",
"start": 0,
"end": 7,
},
{
"entity": "place",
"value": "restaurant",
"start": 22,
"end": 32,
},
{
"entity": "species",
"value": "people",
"start": 55,
"end": 61,
},
])
EntitySynonymEnd().process(example)
def has_changed(entity):
return entity["value"] != example.text[entity["start"]:entity["end"]]
assert example.text == text_initial
changed_entities = filter(has_changed, example.data["entities"])
# Check the unchanged entities match value <-> text[start:end]
assert len(list(changed_entities)) == 3
# Check the changed entities are reverted properly
for initial, entity in zip(initial_entities, changed_entities):
assert raises(KeyError, lambda x: print(x["literal"]), entity)
assert entity["start"] == initial["start"]
assert entity["end"] == initial["end"]
def test_multiple_errors():
instance = _get_instance()
message = Message(text='Ths i a tst mesae')
flagged_tokens = [
{
"offset": 0,
"token": "Ths",
"type": "UnknownToken",
"suggestions": [
{
"suggestion": "This",
"score": 0.825389307284585
}
]
},
{
"offset": 4,
"token": "i",
"type": "UnknownToken",
"suggestions": [
{
"suggestion": "is",
"score": 0.825389307284585
}
]
},
{
"offset": 8,
"token": "tst",
"type": "UnknownToken",
"suggestions": [
{
"suggestion": "test",
"score": 0.825389307284585
},
{
"suggestion": "text",
"score": 0.646529276890009
}
]
},
{
"offset": 12,
"token": "mesae",
"type": "UnknownToken",
"suggestions": [
{
"suggestion": "message",
"score": 0.825389307284585
},
{
"suggestion": "mesa",
"score": 0.761621385590906
}
]
}
]
tokens = instance._get_replacements(flagged_tokens)
assert len(tokens) == len(flagged_tokens)
text = instance._replace(message.text, tokens)
assert text == 'This is a test message'
