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