from typing import Dict
import torch


def create_ner_conditional_masks(id2label: Dict[int, str]) -> torch.Tensor:
    """Create a NER-conditional mask matrix which implies the relations between
    before-tag and after-tag.

    According to the rule of BIO-naming system, it is impossible that `I-Dog` cannot be
    appeard after `B-Dog` or `I-Dog` tags. This function creates the calculable
    relation-based conditional matrix to prevent from generating wrong tags.

    Args:
        id2label: A dictionary which maps class indices to their label names.

    Returns:
        A conditional mask tensor.
    """
    conditional_masks = torch.zeros(len(id2label), len(id2label))
    for i, before in id2label.items():
        for j, after in id2label.items():
            if after == "O" or after.startswith("B-") or after == f"I-{before[2:]}":
                conditional_masks[i, j] = 1.0
    return conditional_masks