def extract_entities(self, message: Message) -> List[Dict[Text, Any]]:
"""Extract entities from the given message using the trained model(s)."""
if self.entity_taggers is None or not message.features_present(
attribute=TEXT,
featurizers=self.component_config.get(FEATURIZERS)):
return []
tokens = message.get(TOKENS_NAMES[TEXT])
crf_tokens = self._convert_to_crf_tokens(message)
predictions: Dict[Text, List[Dict[Text, float]]] = {}
for tag_name, entity_tagger in self.entity_taggers.items():
# use predicted entity tags as features for second level CRFs
include_tag_features = tag_name != ENTITY_ATTRIBUTE_TYPE
if include_tag_features:
self._add_tag_to_crf_token(crf_tokens, predictions)
features = self._crf_tokens_to_features(crf_tokens,
include_tag_features)
predictions[tag_name] = entity_tagger.predict_marginals_single(
features)
# convert predictions into a list of tags and a list of confidences
tags, confidences = self._tag_confidences(tokens, predictions)
return self.convert_predictions_into_entities(
message.get(TEXT), tokens, tags, self.split_entities_config,
confidences)
def test_features_present(
features: Optional[List[Features]],
attribute: Text,
featurizers: List[Text],
expected: bool,
):
message = Message(data={TEXT: "This is a test sentence."}, features=features)
actual = message.features_present(attribute, featurizers)
assert actual == expected