Function align_entity_predictions

rasa/nlu/test.py:1125–1163

Aligns entity predictions to the message tokens. Determines for every token the true label based on the prediction targets and the label assigned by each single extractor. Args: result: entity evaluation result extractors: the entity extractors that should be consid

(
    result: EntityEvaluationResult, extractors: Set[Text]
)

def align_entity_predictions(
    result: EntityEvaluationResult, extractors: Set[Text]
) -> Dict:
    """Aligns entity predictions to the message tokens.

    Determines for every token the true label based on the
    prediction targets and the label assigned by each
    single extractor.

    Args:
        result: entity evaluation result
        extractors: the entity extractors that should be considered

    Returns: dictionary containing the true token labels and token labels
             from the extractors
    """
    true_token_labels = []
    entities_by_extractors: Dict[Text, List] = {
        extractor: [] for extractor in extractors
    }
    for p in result.entity_predictions:
        entities_by_extractors[p[EXTRACTOR]].append(p)
    extractor_labels: Dict[Text, List] = {extractor: [] for extractor in extractors}
    extractor_confidences: Dict[Text, List] = {
        extractor: [] for extractor in extractors
    }
    for t in result.tokens:
        true_token_labels.append(_concat_entity_labels(t, result.entity_targets))
        for extractor, entities in entities_by_extractors.items():
            extracted_labels = _concat_entity_labels(t, entities, {extractor})
            extracted_confidences = _get_entity_confidences(t, entities, {extractor})
            extractor_labels[extractor].append(extracted_labels)
            extractor_confidences[extractor].append(extracted_confidences)

    return {
        "target_labels": true_token_labels,
        "extractor_labels": extractor_labels,
        "confidences": extractor_confidences,
    }


def _concat_entity_labels(