rasa-pro 3.12.0.dev1__py3-none-any.whl

rasa/nlu/utils/bilou_utils.py

@@ -0,0 +1,462 @@
+import logging
+import operator
+from collections import defaultdict, Counter
+from typing import List, Tuple, Text, Optional, Dict, Any, TYPE_CHECKING
+
+from rasa.nlu.constants import (
+    TOKENS_NAMES,
+    BILOU_ENTITIES,
+    BILOU_ENTITIES_GROUP,
+    BILOU_ENTITIES_ROLE,
+)
+from rasa.shared.nlu.constants import (
+    TEXT,
+    ENTITIES,
+    ENTITY_ATTRIBUTE_START,
+    ENTITY_ATTRIBUTE_END,
+    ENTITY_ATTRIBUTE_TYPE,
+    ENTITY_ATTRIBUTE_GROUP,
+    ENTITY_ATTRIBUTE_ROLE,
+    NO_ENTITY_TAG,
+)
+
+if TYPE_CHECKING:
+    from rasa.nlu.tokenizers.tokenizer import Token
+    from rasa.shared.nlu.training_data.training_data import TrainingData
+    from rasa.shared.nlu.training_data.message import Message
+
+logger = logging.getLogger(__name__)
+
+BEGINNING = "B-"
+INSIDE = "I-"
+LAST = "L-"
+UNIT = "U-"
+BILOU_PREFIXES = [BEGINNING, INSIDE, LAST, UNIT]
+
+
+def bilou_prefix_from_tag(tag: Text) -> Optional[Text]:
+    """Returns the BILOU prefix from the given tag.
+
+    Args:
+        tag: the tag
+
+    Returns: the BILOU prefix of the tag
+    """
+    if tag[:2] in BILOU_PREFIXES:
+        return tag[:2]
+    return None
+
+
+def tag_without_prefix(tag: Text) -> Text:
+    """Remove the BILOU prefix from the given tag.
+
+    Args:
+        tag: the tag
+
+    Returns: the tag without the BILOU prefix
+    """
+    if tag[:2] in BILOU_PREFIXES:
+        return tag[2:]
+    return tag
+
+
+def bilou_tags_to_ids(
+    message: "Message",
+    tag_id_dict: Dict[Text, int],
+    tag_name: Text = ENTITY_ATTRIBUTE_TYPE,
+) -> List[int]:
+    """Maps the entity tags of the message to the ids of the provided dict.
+
+    Args:
+        message: the message
+        tag_id_dict: mapping of tags to ids
+        tag_name: tag name of interest
+
+    Returns: a list of tag ids
+    """
+    bilou_key = get_bilou_key_for_tag(tag_name)
+
+    if message.get(bilou_key):
+        _tags = [
+            tag_id_dict[_tag] if _tag in tag_id_dict else tag_id_dict[NO_ENTITY_TAG]
+            for _tag in message.get(bilou_key)
+        ]
+    else:
+        _tags = [tag_id_dict[NO_ENTITY_TAG] for _ in message.get(TOKENS_NAMES[TEXT])]
+
+    return _tags
+
+
+def get_bilou_key_for_tag(tag_name: Text) -> Text:
+    """Get the message key for the BILOU tagging format of the provided tag name.
+
+    Args:
+        tag_name: the tag name
+
+    Returns:
+        the message key to store the BILOU tags
+    """
+    if tag_name == ENTITY_ATTRIBUTE_ROLE:
+        return BILOU_ENTITIES_ROLE
+
+    if tag_name == ENTITY_ATTRIBUTE_GROUP:
+        return BILOU_ENTITIES_GROUP
+
+    return BILOU_ENTITIES
+
+
+def build_tag_id_dict(
+    training_data: "TrainingData", tag_name: Text = ENTITY_ATTRIBUTE_TYPE
+) -> Optional[Dict[Text, int]]:
+    """Create a mapping of unique tags to ids.
+
+    Args:
+        training_data: the training data
+        tag_name: tag name of interest
+
+    Returns: a mapping of tags to ids
+    """
+    bilou_key = get_bilou_key_for_tag(tag_name)
+
+    distinct_tags = set(
+        [
+            tag_without_prefix(e)
+            for example in training_data.nlu_examples
+            if example.get(bilou_key)
+            for e in example.get(bilou_key)
+        ]
+    ) - {NO_ENTITY_TAG}
+
+    if not distinct_tags:
+        return None
+
+    tag_id_dict = {
+        f"{prefix}{tag}": idx_1 * len(BILOU_PREFIXES) + idx_2 + 1
+        for idx_1, tag in enumerate(sorted(distinct_tags))
+        for idx_2, prefix in enumerate(BILOU_PREFIXES)
+    }
+    # NO_ENTITY_TAG corresponds to non-entity which should correspond to 0 index
+    # needed for correct prediction for padding
+    tag_id_dict[NO_ENTITY_TAG] = 0
+
+    return tag_id_dict
+
+
+def apply_bilou_schema(training_data: "TrainingData") -> None:
+    """Get a list of BILOU entity tags and set them on the given messages.
+
+    Args:
+        training_data: the training data
+    """
+    for message in training_data.nlu_examples:
+        apply_bilou_schema_to_message(message)
+
+
+def apply_bilou_schema_to_message(message: "Message") -> None:
+    """Get a list of BILOU entity tags and set them on the given message.
+
+    Args:
+        message: the message
+    """
+    entities = message.get(ENTITIES)
+
+    if not entities:
+        return
+
+    tokens = message.get(TOKENS_NAMES[TEXT])
+
+    for attribute, message_key in [
+        (ENTITY_ATTRIBUTE_TYPE, BILOU_ENTITIES),
+        (ENTITY_ATTRIBUTE_ROLE, BILOU_ENTITIES_ROLE),
+        (ENTITY_ATTRIBUTE_GROUP, BILOU_ENTITIES_GROUP),
+    ]:
+        entities = map_message_entities(message, attribute)
+        output = bilou_tags_from_offsets(tokens, entities)
+        message.set(message_key, output)
+
+
+def map_message_entities(
+    message: "Message", attribute_key: Text = ENTITY_ATTRIBUTE_TYPE
+) -> List[Tuple[int, int, Text]]:
+    """Maps the entities of the given message to their start, end, and tag values.
+
+    Args:
+        message: the message
+        attribute_key: key of tag value to use
+
+    Returns: a list of start, end, and tag value tuples
+    """
+
+    def convert_entity(entity: Dict[Text, Any]) -> Tuple[int, int, Text]:
+        return (
+            entity[ENTITY_ATTRIBUTE_START],
+            entity[ENTITY_ATTRIBUTE_END],
+            entity.get(attribute_key) or NO_ENTITY_TAG,
+        )
+
+    entities = [convert_entity(entity) for entity in message.get(ENTITIES, [])]
+
+    # entities is a list of tuples (start, end, tag value).
+    # filter out all entities with tag value == NO_ENTITY_TAG.
+    tag_value_idx = 2
+    return [entity for entity in entities if entity[tag_value_idx] != NO_ENTITY_TAG]
+
+
+def bilou_tags_from_offsets(
+    tokens: List["Token"], entities: List[Tuple[int, int, Text]]
+) -> List[Text]:
+    """Creates BILOU tags for the given tokens and entities.
+
+    Args:
+        message: The message object.
+        tokens: The list of tokens.
+        entities: The list of start, end, and tag tuples.
+        missing: The tag for missing entities.
+
+    Returns:
+        BILOU tags.
+    """
+    start_pos_to_token_idx = {token.start: i for i, token in enumerate(tokens)}
+    end_pos_to_token_idx = {token.end: i for i, token in enumerate(tokens)}
+
+    bilou = [NO_ENTITY_TAG for _ in tokens]
+
+    _add_bilou_tags_to_entities(
+        bilou, entities, end_pos_to_token_idx, start_pos_to_token_idx
+    )
+
+    return bilou
+
+
+def _add_bilou_tags_to_entities(
+    bilou: List[Text],
+    entities: List[Tuple[int, int, Text]],
+    end_pos_to_token_idx: Dict[int, int],
+    start_pos_to_token_idx: Dict[int, int],
+) -> None:
+    for start_pos, end_pos, label in entities:
+        start_token_idx = start_pos_to_token_idx.get(start_pos)
+        end_token_idx = end_pos_to_token_idx.get(end_pos)
+
+        # Only interested if the tokenization is correct
+        if start_token_idx is not None and end_token_idx is not None:
+            if start_token_idx == end_token_idx:
+                bilou[start_token_idx] = f"{UNIT}{label}"
+            else:
+                bilou[start_token_idx] = f"{BEGINNING}{label}"
+                for i in range(start_token_idx + 1, end_token_idx):
+                    bilou[i] = f"{INSIDE}{label}"
+                bilou[end_token_idx] = f"{LAST}{label}"
+
+
+def ensure_consistent_bilou_tagging(
+    predicted_tags: List[Text], predicted_confidences: List[float]
+) -> Tuple[List[Text], List[float]]:
+    """Ensure predicted tags follow the BILOU tagging schema.
+
+    We assume that starting B- tags are correct. Followed tags that belong to start
+    tag but have a different entity type are updated considering also the confidence
+    values of those tags.
+    For example, B-a I-b L-a is updated to B-a I-a L-a and B-a I-a O is changed to
+    B-a L-a.
+
+    Args:
+        predicted_tags: predicted tags
+        predicted_confidences: predicted confidences
+
+    Return:
+        List of tags.
+        List of confidences.
+    """
+    for idx, predicted_tag in enumerate(predicted_tags):
+        prefix = bilou_prefix_from_tag(predicted_tag)
+        tag = tag_without_prefix(predicted_tag)
+
+        if prefix == BEGINNING:
+            last_idx = _find_bilou_end(idx, predicted_tags)
+
+            relevant_confidences = predicted_confidences[idx : last_idx + 1]
+            relevant_tags = [
+                tag_without_prefix(tag) for tag in predicted_tags[idx : last_idx + 1]
+            ]
+
+            # if not all tags are the same, for example, B-person I-person L-location
+            # we need to check what tag we should use depending on the confidence
+            # values and update the tags and confidences accordingly
+            if not all(relevant_tags[0] == tag for tag in relevant_tags):
+                # decide which tag this entity should use
+                tag, tag_score = _tag_to_use(relevant_tags, relevant_confidences)
+
+                logger.debug(
+                    f"Using tag '{tag}' for entity with mixed tag labels "
+                    f"(original tags: {predicted_tags[idx : last_idx + 1]}, "
+                    f"(original confidences: "
+                    f"{predicted_confidences[idx : last_idx + 1]})."
+                )
+
+                # all tags that change get the score of that tag assigned
+                predicted_confidences = _update_confidences(
+                    predicted_confidences, predicted_tags, tag, tag_score, idx, last_idx
+                )
+
+            # ensure correct BILOU annotations
+            if last_idx == idx:
+                predicted_tags[idx] = f"{UNIT}{tag}"
+            elif last_idx - idx == 1:
+                predicted_tags[idx] = f"{BEGINNING}{tag}"
+                predicted_tags[last_idx] = f"{LAST}{tag}"
+            else:
+                predicted_tags[idx] = f"{BEGINNING}{tag}"
+                predicted_tags[last_idx] = f"{LAST}{tag}"
+                for i in range(idx + 1, last_idx):
+                    predicted_tags[i] = f"{INSIDE}{tag}"
+
+    return predicted_tags, predicted_confidences
+
+
+def _tag_to_use(
+    relevant_tags: List[Text], relevant_confidences: List[float]
+) -> Tuple[Text, float]:
+    """Decide what tag to use according to the following metric:
+
+    Calculate the average confidence per tag.
+    Calculate the percentage of tokens assigned to a tag within the entity per tag.
+    The harmonic mean of those two metrics is the score for the tag.
+    The tag with the highest score is taken as the tag for the entity.
+
+    Args:
+        relevant_tags: The tags of the entity.
+        relevant_confidences: The confidence values.
+
+    Returns:
+        The tag to use. The score of that tag.
+    """
+    # Calculate the average confidence per tag.
+    avg_confidence_per_tag = _avg_confidence_per_tag(
+        relevant_tags, relevant_confidences
+    )
+    # Calculate the percentage of tokens assigned to a tag per tag.
+    tag_counts = Counter(relevant_tags)
+    token_percentage_per_tag: Dict[Text, float] = {}
+    for tag, count in tag_counts.items():
+        token_percentage_per_tag[tag] = round(count / len(relevant_tags), 2)
+
+    # Calculate the harmonic mean between the two metrics per tag.
+    score_per_tag = {}
+    for tag, token_percentage in token_percentage_per_tag.items():
+        avg_confidence = avg_confidence_per_tag[tag]
+        score_per_tag[tag] = (
+            2
+            * (avg_confidence * token_percentage)
+            / (avg_confidence + token_percentage)
+        )
+
+    # Take the tag with the highest score as the tag for the entity
+    tag, score = max(score_per_tag.items(), key=operator.itemgetter(1))
+
+    return tag, score
+
+
+def _update_confidences(
+    predicted_confidences: List[float],
+    predicted_tags: List[Text],
+    tag: Text,
+    score: float,
+    idx: int,
+    last_idx: int,
+) -> List[float]:
+    """Update the confidence values.
+
+    Set the confidence value of a tag to score value if the predicated
+    tag changed.
+
+    Args:
+        predicted_confidences: The list of predicted confidences.
+        predicted_tags: The list of predicted tags.
+        tag: The tag of the entity.
+        score: The score value of that tag.
+        idx: The start index of the entity.
+        last_idx: The end index of the entity.
+
+    Returns:
+        The updated list of confidences.
+    """
+    for i in range(idx, last_idx + 1):
+        predicted_confidences[i] = (
+            round(score, 2)
+            if tag_without_prefix(predicted_tags[i]) != tag
+            else predicted_confidences[i]
+        )
+    return predicted_confidences
+
+
+def _avg_confidence_per_tag(
+    relevant_tags: List[Text], relevant_confidences: List[float]
+) -> Dict[Text, float]:
+    confidences_per_tag = defaultdict(list)
+
+    for tag, confidence in zip(relevant_tags, relevant_confidences):
+        confidences_per_tag[tag].append(confidence)
+
+    avg_confidence_per_tag = {}
+    for tag, confidences in confidences_per_tag.items():
+        avg_confidence_per_tag[tag] = round(sum(confidences) / len(confidences), 2)
+
+    return avg_confidence_per_tag
+
+
+def _find_bilou_end(start_idx: int, predicted_tags: List[Text]) -> int:
+    """Find the last index of the entity.
+
+    The start index is pointing to a B- tag. The entity is closed as soon as we find
+    a L- tag or a O tag.
+
+    Args:
+        start_idx: The start index of the entity
+        predicted_tags: The list of predicted tags
+
+    Returns:
+        The end index of the entity
+    """
+    current_idx = start_idx + 1
+    finished = False
+    start_tag = tag_without_prefix(predicted_tags[start_idx])
+
+    while not finished:
+        if current_idx >= len(predicted_tags):
+            logger.debug(
+                "Inconsistent BILOU tagging found, B- tag not closed by L- tag, "
+                "i.e [B-a, I-a, O] instead of [B-a, L-a, O].\n"
+                "Assuming last tag is L- instead of I-."
+            )
+            current_idx -= 1
+            break
+
+        current_label = predicted_tags[current_idx]
+        prefix = bilou_prefix_from_tag(current_label)
+        tag = tag_without_prefix(current_label)
+
+        if tag != start_tag:
+            # words are not tagged the same entity class
+            logger.debug(
+                "Inconsistent BILOU tagging found, B- tag, L- tag pair encloses "
+                "multiple entity classes.i.e. [B-a, I-b, L-a] instead of "
+                "[B-a, I-a, L-a].\nAssuming B- class is correct."
+            )
+
+        if prefix == LAST:
+            finished = True
+        elif prefix == INSIDE:
+            # middle part of the entity
+            current_idx += 1
+        else:
+            # entity not closed by an L- tag
+            finished = True
+            current_idx -= 1
+            logger.debug(
+                "Inconsistent BILOU tagging found, B- tag not closed by L- tag, "
+                "i.e [B-a, I-a, O] instead of [B-a, L-a, O].\n"
+                "Assuming last tag is L- instead of I-."
+            )
+
+    return current_idx

rasa/nlu/utils/hugging_face/registry.py

@@ -0,0 +1,108 @@
+import logging
+from typing import Dict, Text, Type
+
+# Explicitly set logging level for this module before any import
+# because otherwise it logs tensorflow/pytorch versions
+logging.getLogger("transformers.file_utils").setLevel(logging.WARNING)
+
+from transformers import (  # noqa: E402
+    TFPreTrainedModel,
+    TFBertModel,
+    TFOpenAIGPTModel,
+    TFGPT2Model,
+    TFXLNetModel,
+    # TFXLMModel,
+    TFDistilBertModel,
+    TFRobertaModel,
+    TFCamembertModel,
+    PreTrainedTokenizer,
+    BertTokenizer,
+    OpenAIGPTTokenizer,
+    GPT2Tokenizer,
+    XLNetTokenizer,
+    # XLMTokenizer,
+    DistilBertTokenizer,
+    RobertaTokenizer,
+    CamembertTokenizer,
+)
+from rasa.nlu.utils.hugging_face.transformers_pre_post_processors import (  # noqa: E402
+    bert_tokens_pre_processor,
+    gpt_tokens_pre_processor,
+    xlnet_tokens_pre_processor,
+    roberta_tokens_pre_processor,
+    bert_embeddings_post_processor,
+    gpt_embeddings_post_processor,
+    xlnet_embeddings_post_processor,
+    roberta_embeddings_post_processor,
+    bert_tokens_cleaner,
+    openaigpt_tokens_cleaner,
+    gpt2_tokens_cleaner,
+    xlnet_tokens_cleaner,
+    camembert_tokens_pre_processor,
+)
+
+
+model_class_dict: Dict[Text, Type[TFPreTrainedModel]] = {
+    "bert": TFBertModel,
+    "gpt": TFOpenAIGPTModel,
+    "gpt2": TFGPT2Model,
+    "xlnet": TFXLNetModel,
+    # "xlm": TFXLMModel, # Currently doesn't work because of a bug in transformers
+    # library https://github.com/huggingface/transformers/issues/2729
+    "distilbert": TFDistilBertModel,
+    "roberta": TFRobertaModel,
+    "camembert": TFCamembertModel,
+}
+model_tokenizer_dict: Dict[Text, Type[PreTrainedTokenizer]] = {
+    "bert": BertTokenizer,
+    "gpt": OpenAIGPTTokenizer,
+    "gpt2": GPT2Tokenizer,
+    "xlnet": XLNetTokenizer,
+    # "xlm": XLMTokenizer,
+    "distilbert": DistilBertTokenizer,
+    "roberta": RobertaTokenizer,
+    "camembert": CamembertTokenizer,
+}
+model_weights_defaults = {
+    "bert": "rasa/LaBSE",
+    "gpt": "openai-gpt",
+    "gpt2": "gpt2",
+    "xlnet": "xlnet-base-cased",
+    # "xlm": "xlm-mlm-enfr-1024",
+    "distilbert": "distilbert-base-uncased",
+    "roberta": "roberta-base",
+    "camembert": "camembert-base",
+}
+
+model_special_tokens_pre_processors = {
+    "bert": bert_tokens_pre_processor,
+    "gpt": gpt_tokens_pre_processor,
+    "gpt2": gpt_tokens_pre_processor,
+    "xlnet": xlnet_tokens_pre_processor,
+    # "xlm": xlm_tokens_pre_processor,
+    "distilbert": bert_tokens_pre_processor,
+    "roberta": roberta_tokens_pre_processor,
+    "camembert": camembert_tokens_pre_processor,
+}
+
+model_tokens_cleaners = {
+    "bert": bert_tokens_cleaner,
+    "gpt": openaigpt_tokens_cleaner,
+    "gpt2": gpt2_tokens_cleaner,
+    "xlnet": xlnet_tokens_cleaner,
+    # "xlm": xlm_tokens_pre_processor,
+    "distilbert": bert_tokens_cleaner,  # uses the same as BERT
+    "roberta": gpt2_tokens_cleaner,  # Uses the same as GPT2
+    "camembert": xlnet_tokens_cleaner,  # Removing underscores _
+}
+
+model_embeddings_post_processors = {
+    "bert": bert_embeddings_post_processor,
+    "gpt": gpt_embeddings_post_processor,
+    "gpt2": gpt_embeddings_post_processor,
+    "xlnet": xlnet_embeddings_post_processor,
+    # "xlm": xlm_embeddings_post_processor,
+    "distilbert": bert_embeddings_post_processor,
+    "roberta": roberta_embeddings_post_processor,
+    "camembert": roberta_embeddings_post_processor,
+}