rasahq--rasa/tests/nlu/extractors/test_extractor.py

from typing import Any, Text, Dict, List

import pytest
from rasa.shared.constants import LATEST_TRAINING_DATA_FORMAT_VERSION

from rasa.shared.nlu.constants import TEXT, SPLIT_ENTITIES_BY_COMMA
from rasa.shared.nlu.training_data.message import Message
from rasa.nlu.extractors.extractor import EntityExtractorMixin
from rasa.nlu.tokenizers.whitespace_tokenizer import WhitespaceTokenizer
from rasa.shared.nlu.training_data.formats.rasa_yaml import RasaYAMLReader


@pytest.mark.parametrize(
    "text, tags, confidences, expected_entities",
    [
        (
            "I am flying from San Fransisco to Amsterdam",
            {
                "entity": ["O", "O", "O", "O", "city", "city", "O", "city"],
                "role": ["O", "O", "O", "O", "from", "from", "O", "to"],
            },
            {
                "entity": [1.0, 1.0, 1.0, 1.0, 0.98, 0.78, 1.0, 0.89],
                "role": [1.0, 1.0, 1.0, 1.0, 0.98, 0.78, 1.0, 0.89],
            },
            [
                {
                    "entity": "city",
                    "start": 17,
                    "end": 30,
                    "value": "San Fransisco",
                    "role": "from",
                    "confidence_entity": 0.78,
                    "confidence_role": 0.78,
                },
                {
                    "entity": "city",
                    "start": 34,
                    "end": 43,
                    "value": "Amsterdam",
                    "role": "to",
                    "confidence_entity": 0.89,
                    "confidence_role": 0.89,
                },
            ],
        ),
        (
            "I am flying from San Fransisco to Amsterdam",
            {
                "entity": ["O", "O", "O", "O", "B-city", "L-city", "O", "U-city"],
                "role": ["O", "O", "O", "O", "B-from", "L-from", "O", "U-to"],
            },
            {
                "entity": [1.0, 1.0, 1.0, 1.0, 0.98, 0.78, 1.0, 0.89],
                "role": [1.0, 1.0, 1.0, 1.0, 0.98, 0.78, 1.0, 0.89],
            },
            [
                {
                    "entity": "city",
                    "start": 17,
                    "end": 30,
                    "value": "San Fransisco",
                    "role": "from",
                    "confidence_entity": 0.78,
                    "confidence_role": 0.78,
                },
                {
                    "entity": "city",
                    "start": 34,
                    "end": 43,
                    "value": "Amsterdam",
                    "role": "to",
                    "confidence_entity": 0.89,
                    "confidence_role": 0.89,
                },
            ],
        ),
        (
            "I am flying from San Fransisco to Amsterdam",
            {
                "entity": ["O", "O", "O", "O", "city", "city", "O", "city"],
                "group": ["O", "O", "O", "O", "1", "1", "O", "1"],
            },
            None,
            [
                {
                    "entity": "city",
                    "start": 17,
                    "end": 30,
                    "value": "San Fransisco",
                    "group": "1",
                },
                {
                    "entity": "city",
                    "start": 34,
                    "end": 43,
                    "value": "Amsterdam",
                    "group": "1",
                },
            ],
        ),
        (
            "Amsterdam",
            {"entity": ["city"], "role": ["O"], "group": ["O"]},
            None,
            [{"entity": "city", "start": 0, "end": 9, "value": "Amsterdam"}],
        ),
        (
            "New-York",
            {"entity": ["city", "city"], "role": ["O", "O"], "group": ["O", "O"]},
            None,
            [{"entity": "city", "start": 0, "end": 8, "value": "New-York"}],
        ),
        (
            "Amsterdam, Berlin, and London",
            {
                "entity": ["city", "city", "O", "city"],
                "role": ["O", "O", "O", "O"],
                "group": ["O", "O", "O", "O"],
            },
            None,
            [
                {"entity": "city", "start": 0, "end": 9, "value": "Amsterdam"},
                {"entity": "city", "start": 11, "end": 17, "value": "Berlin"},
                {"entity": "city", "start": 23, "end": 29, "value": "London"},
            ],
        ),
        (
            "Amsterdam Berlin and London",
            {
                "entity": ["U-city", "U-city", "O", "U-city"],
                "role": ["O", "O", "O", "O"],
                "group": ["O", "O", "O", "O"],
            },
            None,
            [
                {"entity": "city", "start": 0, "end": 9, "value": "Amsterdam"},
                {"entity": "city", "start": 10, "end": 16, "value": "Berlin"},
                {"entity": "city", "start": 21, "end": 27, "value": "London"},
            ],
        ),
        (
            "San Fransisco Amsterdam, London",
            {
                "entity": ["B-city", "L-city", "U-city", "U-city"],
                "role": ["O", "O", "O", "O"],
                "group": ["O", "O", "O", "O"],
            },
            None,
            [
                {"entity": "city", "start": 0, "end": 13, "value": "San Fransisco"},
                {"entity": "city", "start": 14, "end": 23, "value": "Amsterdam"},
                {"entity": "city", "start": 25, "end": 31, "value": "London"},
            ],
        ),
        (
            "New York City Los Angeles and San Diego",
            {
                "entity": [
                    "B-city",
                    "I-city",
                    "L-city",
                    "B-city",
                    "L-city",
                    "O",
                    "B-city",
                    "L-city",
                ],
                "role": ["O", "O", "O", "O", "O", "O", "O", "O"],
                "group": ["O", "O", "O", "O", "O", "O", "O", "O"],
            },
            None,
            [
                {"entity": "city", "start": 0, "end": 13, "value": "New York City"},
                {"entity": "city", "start": 14, "end": 25, "value": "Los Angeles"},
                {"entity": "city", "start": 30, "end": 39, "value": "San Diego"},
            ],
        ),
        (
            "Berlin weather",
            {"entity": ["I-city", "O"], "role": ["O", "O"], "group": ["O", "O"]},
            None,
            [{"entity": "city", "start": 0, "end": 6, "value": "Berlin"}],
        ),
        (
            "New-York",
            {"entity": ["city", "city"], "role": ["O", "O"], "group": ["O", "O"]},
            None,
            [{"entity": "city", "start": 0, "end": 8, "value": "New-York"}],
        ),
        (
            "Amsterdam, Berlin, and London",
            {
                "entity": ["city", "city", "O", "city"],
                "role": ["O", "O", "O", "O"],
                "group": ["O", "O", "O", "O"],
            },
            None,
            [
                {"entity": "city", "start": 0, "end": 9, "value": "Amsterdam"},
                {"entity": "city", "start": 11, "end": 17, "value": "Berlin"},
                {"entity": "city", "start": 23, "end": 29, "value": "London"},
            ],
        ),
        (
            "Amsterdam Berlin and London",
            {
                "entity": ["city", "city", "O", "city"],
                "role": ["O", "O", "O", "O"],
                "group": ["O", "O", "O", "O"],
            },
            None,
            [
                {"entity": "city", "start": 0, "end": 16, "value": "Amsterdam Berlin"},
                {"entity": "city", "start": 21, "end": 27, "value": "London"},
            ],
        ),
    ],
)
def test_convert_tags_to_entities(
    text: Text,
    tags: Dict[Text, List[Text]],
    confidences: Dict[Text, List[float]],
    expected_entities: List[Dict[Text, Any]],
    whitespace_tokenizer: WhitespaceTokenizer,
):
    extractor = EntityExtractorMixin()

    message = Message(data={TEXT: text})
    tokens = whitespace_tokenizer.tokenize(message, TEXT)

    split_entities_config = {SPLIT_ENTITIES_BY_COMMA: True}
    actual_entities = extractor.convert_predictions_into_entities(
        text, tokens, tags, split_entities_config, confidences
    )
    assert actual_entities == expected_entities


@pytest.mark.parametrize(
    "text, tags, confidences, expected_entities",
    [
        (
            "I live at 22 Powderhall Rd., EH7 4GB, Edinburgh, UK",
            {
                "entity": [
                    "O",
                    "O",
                    "O",
                    "address",
                    "address",
                    "address",
                    "address",
                    "address",
                    "address",
                    "address",
                ]
            },
            {"entity": [1.0, 1.0, 1.0, 1.0, 0.98, 0.78, 1.0, 0.89, 1.0, 1.0, 1.0]},
            [
                {
                    "entity": "address",
                    "start": 10,
                    "end": 51,
                    "value": "22 Powderhall Rd., EH7 4GB, Edinburgh, UK",
                    "confidence_entity": 0.78,
                }
            ],
        ),
        (
            "The address is Schönhauser Allee 175, 10119 Berlin, DE",
            {
                "entity": [
                    "O",
                    "O",
                    "O",
                    "address",
                    "address",
                    "address",
                    "address",
                    "address",
                    "address",
                ]
            },
            {"entity": [1.0, 1.0, 1.0, 1.0, 1.0, 0.67, 0.77, 1.0, 0.98]},
            [
                {
                    "entity": "address",
                    "start": 15,
                    "end": 54,
                    "value": "Schönhauser Allee 175, 10119 Berlin, DE",
                    "confidence_entity": 0.67,
                }
            ],
        ),
        (
            "We need to get more of tofu, cauliflower, avocado",
            {
                "entity": [
                    "O",
                    "O",
                    "O",
                    "O",
                    "O",
                    "O",
                    "ingredient",
                    "ingredient",
                    "ingredient",
                ]
            },
            {"entity": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.9, 0.78]},
            [
                {
                    "entity": "ingredient",
                    "start": 23,
                    "end": 27,
                    "value": "tofu",
                    "confidence_entity": 1.0,
                },
                {
                    "entity": "ingredient",
                    "start": 29,
                    "end": 40,
                    "value": "cauliflower",
                    "confidence_entity": 0.9,
                },
                {
                    "entity": "ingredient",
                    "start": 42,
                    "end": 49,
                    "value": "avocado",
                    "confidence_entity": 0.78,
                },
            ],
        ),
        (
            "So the list of drinks to get is coffee, Club Mate, Ottakringer",
            {
                "entity": [
                    "O",
                    "O",
                    "O",
                    "O",
                    "O",
                    "O",
                    "O",
                    "O",
                    "beverage",
                    "beverage",
                    "beverage",
                    "beverage",
                ]
            },
            {
                "entity": [
                    1.0,
                    1.0,
                    1.0,
                    1.0,
                    1.0,
                    1.0,
                    1.0,
                    1.0,
                    0.69,
                    0.88,
                    0.84,
                    0.79,
                ]
            },
            [
                {
                    "entity": "beverage",
                    "start": 32,
                    "end": 38,
                    "value": "coffee",
                    "confidence_entity": 0.69,
                },
                {
                    "entity": "beverage",
                    "start": 40,
                    "end": 49,
                    "value": "Club Mate",
                    "confidence_entity": 0.84,
                },
                {
                    "entity": "beverage",
                    "start": 51,
                    "end": 62,
                    "value": "Ottakringer",
                    "confidence_entity": 0.79,
                },
            ],
        ),
    ],
)
def test_split_entities_by_comma(
    text: Text,
    tags: Dict[Text, List[Text]],
    confidences: Dict[Text, List[float]],
    expected_entities: List[Dict[Text, Any]],
    whitespace_tokenizer: WhitespaceTokenizer,
):
    extractor = EntityExtractorMixin()

    message = Message(data={TEXT: text})
    tokens = whitespace_tokenizer.tokenize(message, TEXT)

    split_entities_config = {
        SPLIT_ENTITIES_BY_COMMA: True,
        "address": False,
        "ingredient": True,
    }
    actual_entities = extractor.convert_predictions_into_entities(
        text, tokens, tags, split_entities_config, confidences
    )

    assert actual_entities == expected_entities


@pytest.mark.parametrize(
    "text, warnings",
    [
        (
            f'version: "{LATEST_TRAINING_DATA_FORMAT_VERSION}"\n'
            "nlu:\n"
            "- intent: test\n"
            "  examples: |\n"
            "    - I want to fly from [Berlin](location) to [ London](location)\n",
            1,
        ),
        (
            f'version: "{LATEST_TRAINING_DATA_FORMAT_VERSION}"\n'
            "nlu:\n"
            "- intent: test\n"
            "  examples: |\n"
            "    - I want to fly from [Berlin ](location) to [London](location)\n",
            1,
        ),
        (
            f'version: "{LATEST_TRAINING_DATA_FORMAT_VERSION}"\n'
            "nlu:\n"
            "- intent: test\n"
            "  examples: |\n"
            "    - I want to fly from [Berlin](location) to [London.](location)\n"
            "    - I have nothing to say.\n",
            1,
        ),
        (
            f'version: "{LATEST_TRAINING_DATA_FORMAT_VERSION}"\n'
            "nlu:\n"
            "- intent: test\n"
            "  examples: |\n"
            "    - I have nothing to say.\n"
            "    - I want to fly from [Berlin](location) to[San Fransisco](location)\n",
            1,
        ),
        (
            f'version: "{LATEST_TRAINING_DATA_FORMAT_VERSION}"\n'
            "nlu:\n"
            "- intent: test\n"
            "  examples: |\n"
            "    - I want to fly from [Berlin](location) to[San Fransisco](location)\n"
            "    - Book a flight from [London](location) to [Paris.](location)\n",
            2,
        ),
    ],
)
def test_check_correct_entity_annotations(
    text: Text, warnings: int, whitespace_tokenizer: WhitespaceTokenizer
):
    reader = RasaYAMLReader()

    training_data = reader.reads(text)
    whitespace_tokenizer.process_training_data(training_data)

    with pytest.warns(UserWarning) as record:
        EntityExtractorMixin.check_correct_entity_annotations(training_data)

    assert len(record) == warnings
    assert all(
        [excerpt in record[0].message.args[0]]
        for excerpt in ["Misaligned entity annotation in sentence"]
    )