<reponame>bbieniek/deidentify<filename>deidentify/methods/tagging_utils.py
"""Utility methods to convert between standoff and BIO format.
"""
import warnings
from collections import defaultdict, namedtuple
from typing import List, Tuple

import spacy
from loguru import logger

try:
    from spacy.gold import biluo_tags_from_offsets, offsets_from_biluo_tags
except ModuleNotFoundError:
    # spacy>=3
    from spacy.training.iob_utils import biluo_tags_from_offsets, offsets_from_biluo_tags

from tqdm import tqdm

from deidentify.base import Annotation, Document
from deidentify.tokenizer import Tokenizer

Token = namedtuple('Token', ['text', 'pos_tag', 'label', 'ner_tag'])
ParsedDoc = namedtuple('ParsedDoc', ['spacy_doc', 'name', 'text'])

# Silence spaCy warning regarding misaligned entity boundaries. It will show up multiple times
# because the message changes with the input text.
# More info on the warning: https://github.com/explosion/spaCy/issues/5727
warnings.filterwarnings('ignore', message=r'.*W030.*')


def standoff_to_sents(docs: List[Document],
                      tokenizer: Tokenizer,
                      verbose=False) -> Tuple[List[List[Token]], List[ParsedDoc]]:
    """Convert corpus into list of BIO tagged sentences.

    Each document is parsed using the spaCy tokenizer and segmented into sentences. Afterwards, each
    sentence is represented as a list of BIO tagged tokens.

    Parameters
    ----------
    docs : List[Document]
        The corpus documents.
    tokenizer : Tokenizer
        A tokenizer instance that parsed the documents using spaCy.

    Returns
    -------
    sents : List[List[Token]]
        List of sentences, where each sentence is a list of BIO tagged tokens.
    sents_docs : List[ParsedDoc]
        The parsed documents where the sentence originates from.
    """
    sents = []
    sents_docs = []

    for doc in tqdm(docs, disable=not verbose, desc='Tokenize documents'):
        parsed_doc = tokenizer.parse_text(doc.text)

        bio_tags = _doc_to_bio(parsed_doc, doc.annotations)

        for sent in parsed_doc.sents:
            sent_tokens = [Token(text=token.text,
                                 pos_tag=token.pos_,
                                 label=bio_tags[token.i],
                                 ner_tag=token.ent_type_)
                           for token in sent]

            # merge sentences where entities cross sentence boundaries
            if sent_tokens[0].label.startswith('I-') or sent_tokens[0].label.startswith('L-'):
                sents[-1].extend(sent_tokens)
            else:
                sents.append(sent_tokens)
                sents_docs.append(ParsedDoc(spacy_doc=parsed_doc, name=doc.name, text=doc.text))

    _validate_labels(sents_docs, sents)
    return sents, sents_docs


def sents_to_standoff(sentence_tags: List[List[str]], docs: List[ParsedDoc]) -> List[Document]:
    """Convert a BIO tagged documents to standoff annotated documents.

    Parameters
    ----------
    sentence_tags : List[List[str]]
        List of sentences of BIO tagged tokens.
    docs : List[ParsedDoc]
        The documents corresponding to each sentence.

    Returns
    -------
    annotated_docs : List[Document]
        The documents with annotated entities in standoff format.

    """
    tags_by_doc = _group_sentences(sentence_tags, docs)

    annotated_docs = []
    for doc, tags in tags_by_doc:
        try:
            annotated_docs.append(Document(
                name=doc.name,
                text=doc.text,
                annotations=_bio_to_standoff(tags, doc.spacy_doc)
            ))
        except Exception as e:
            logger.warning('Could not convert document to standoff {}\n tags = {}\n{}'
                           .format(doc.name, tags, e))

            annotated_docs.append(Document(
                name=doc.name,
                text=doc.text,
                annotations=[]
            ))

    return annotated_docs


def _group_sentences(sentence_tags: List[List[str]],
                     docs: List[ParsedDoc]) -> List[Tuple[ParsedDoc, List[str]]]:
    """Group BIO tagged sentences by document (i.e., merge list of sentences to a single list per
    document).

    This function is used to convert sentence level BIO annotations to standoff format.

    Example:
    ```
    sentence_tags = [['B', 'I', 'O'], ['O', 'O'], ['O', 'O']]
    docs = [d1, d1, d2]

    _group_sentences(sentence_tags, docs)
    -> [['B', 'I', 'O', 'O', 'O'], ['O', 'O']]
    -> [d1, d2]
    ```

    Parameters
    ----------
    sentence_tags : List[List[str]]
        List of sentences. Each sentence contains a list of BIO tags.
    docs : List[ParsedDoc]
        List of documents corresponding to the sentences.

    Returns
    -------
    grouped_sentences : List[Tuple(ParsedDoc, List[str])]
        The sentences grouped by document.

    """
    grouped = defaultdict(list)
    for doc, sentence in zip(docs, sentence_tags):
        grouped[doc].extend(sentence)

    return list(zip(grouped.keys(), grouped.values()))


def _end_of_chunk(next_tag: str) -> bool:
    return not next_tag or next_tag.startswith('O') or next_tag.startswith('B')


def _bio_to_biluo(bio_tags: List[str]) -> List[str]:
    """Convert sentece-level BIO tagging to BILUO.
    """
    biluo_tags = []
    for i in range(len(bio_tags)):

        current_tag = bio_tags[i]
        if i + 1 < len(bio_tags):
            next_tag = bio_tags[i + 1]
        else:
            next_tag = None

        if current_tag.startswith('B-') and _end_of_chunk(next_tag):
            biluo_tags.append('U-' + current_tag[2:])
        elif current_tag.startswith('I-') and _end_of_chunk(next_tag):
            biluo_tags.append('L-' + current_tag[2:])
        else:
            biluo_tags.append(current_tag)

    return biluo_tags


def _bio_to_standoff(bio_tags: List[str], spacy_doc: spacy.tokens.Doc) -> List[Annotation]:
    """Convert BIO tagged document to annotations in standoff format.

    The original spaCy document is used to recreate correct entity offsets.

    Parameters
    ----------
    bio_tags : List[str]
        A BIO tagged sentence. `len(bio_tags) == len(spacy_doc)` has to hold.
    spacy_doc : spacy.tokens.Doc
        The spaCy doc corresponding to the BIO tags.

    Returns
    -------
    List[Annotation]
        The standoff annotations.

    """
    bio_tags = fix_dangling_entities(bio_tags)
    biluo_tags = _bio_to_biluo(bio_tags)
    offsets = offsets_from_biluo_tags(spacy_doc, biluo_tags)

    annotations = []
    for i, offset in enumerate(offsets):
        annotations.append(Annotation(
            text=spacy_doc.char_span(offset[0], offset[1]).text,
            start=offset[0],
            end=offset[1],
            tag=offset[2],
            ann_id='T{}'.format(i),
        ))

    return annotations


def _doc_to_bio(parsed_doc: spacy.tokens.Doc, annotations: List[Annotation]):
    entities = [(int(ann.start), int(ann.end), ann.tag) for ann in annotations]
    biluo_tags = biluo_tags_from_offsets(parsed_doc, entities)

    biluo_to_bio = {
        'B-': 'B-',
        'I-': 'I-',
        'L-': 'I-',
        'U-': 'B-',
    }

    tags = []
    for tag in biluo_tags:
        if tag == "O":
            tags.append('O')
        elif tag == '-':
            # Returned by spacy if token boundaries mismatch entity boundaries.
            # https://spacy.io/api/goldparse#biluo_tags_from_offsets
            tags.append('O')
            warnings.warn(
                'Some entities could not be aligned in the text. Use `spacy.gold.biluo_tags_from_offsets(nlp.make_doc(text), entities)` to check the alignment.',
                UserWarning
            )
        else:
            tags.append(biluo_to_bio[tag[0:2]] + tag[2:])

    return tags


def _validate_labels(docs: List[ParsedDoc], sents: List[List[Token]]):
    """Perform sanity checks on BIO tagged sentences.

    * A sentence has to start with either a 'B' or 'O' token.
    * If a token is within an entity (i.e., tagged as 'I'), the preceeding token has to be tagged
    as 'I' or 'B'.

    """
    for doc, sent in zip(docs, sents):
        try:
            assert sent[0].label == 'O' or sent[0].label.startswith('B-')
            for i, token in enumerate(sent[1:]):
                if token.label.startswith('I-'):
                    assert sent[i].label.startswith('B-') or sent[i].label.startswith('I-')
        except:
            logger.warning('Invalid tagging {}, sent={}'.format(doc.name, sent))


def _inside_entity(previous_tag, current_tag):
    return previous_tag != 'O' and previous_tag[2:] == current_tag[2:]


def fix_dangling_entities(tags: List[str]) -> List[str]:
    prev_tag = '<OOB>'
    result = []

    for tag in tags:
        if tag.startswith('I-') and not _inside_entity(prev_tag, tag):
            result.append('B-' + tag[2:])
        else:
            result.append(tag)
        prev_tag = tag

    return result

"""Utility methods to convert between standoff and BIO format.
"""
import warnings
from collections import defaultdict, namedtuple
from typing import List, Tuple

import spacy
from loguru import logger

try:
    from spacy.gold import biluo_tags_from_offsets, offsets_from_biluo_tags
except ModuleNotFoundError:
    # spacy>=3
    from spacy.training.iob_utils import biluo_tags_from_offsets, offsets_from_biluo_tags

from tqdm import tqdm

from deidentify.base import Annotation, Document
from deidentify.tokenizer import Tokenizer

Token = namedtuple('Token', ['text', 'pos_tag', 'label', 'ner_tag'])
ParsedDoc = namedtuple('ParsedDoc', ['spacy_doc', 'name', 'text'])

# Silence spaCy warning regarding misaligned entity boundaries. It will show up multiple times
# because the message changes with the input text.
# More info on the warning: https://github.com/explosion/spaCy/issues/5727
warnings.filterwarnings('ignore', message=r'.*W030.*')


def standoff_to_sents(docs: List[Document],
                      tokenizer: Tokenizer,
                      verbose=False) -> Tuple[List[List[Token]], List[ParsedDoc]]:
    """Convert corpus into list of BIO tagged sentences.

    Each document is parsed using the spaCy tokenizer and segmented into sentences. Afterwards, each
    sentence is represented as a list of BIO tagged tokens.

    Parameters
    ----------
    docs : List[Document]
        The corpus documents.
    tokenizer : Tokenizer
        A tokenizer instance that parsed the documents using spaCy.

    Returns
    -------
    sents : List[List[Token]]
        List of sentences, where each sentence is a list of BIO tagged tokens.
    sents_docs : List[ParsedDoc]
        The parsed documents where the sentence originates from.
    """
    sents = []
    sents_docs = []

    for doc in tqdm(docs, disable=not verbose, desc='Tokenize documents'):
        parsed_doc = tokenizer.parse_text(doc.text)

        bio_tags = _doc_to_bio(parsed_doc, doc.annotations)

        for sent in parsed_doc.sents:
            sent_tokens = [Token(text=token.text,
                                 pos_tag=token.pos_,
                                 label=bio_tags[token.i],
                                 ner_tag=token.ent_type_)
                           for token in sent]

            # merge sentences where entities cross sentence boundaries
            if sent_tokens[0].label.startswith('I-') or sent_tokens[0].label.startswith('L-'):
                sents[-1].extend(sent_tokens)
            else:
                sents.append(sent_tokens)
                sents_docs.append(ParsedDoc(spacy_doc=parsed_doc, name=doc.name, text=doc.text))

    _validate_labels(sents_docs, sents)
    return sents, sents_docs


def sents_to_standoff(sentence_tags: List[List[str]], docs: List[ParsedDoc]) -> List[Document]:
    """Convert a BIO tagged documents to standoff annotated documents.

    Parameters
    ----------
    sentence_tags : List[List[str]]
        List of sentences of BIO tagged tokens.
    docs : List[ParsedDoc]
        The documents corresponding to each sentence.

    Returns
    -------
    annotated_docs : List[Document]
        The documents with annotated entities in standoff format.

    """
    tags_by_doc = _group_sentences(sentence_tags, docs)

    annotated_docs = []
    for doc, tags in tags_by_doc:
        try:
            annotated_docs.append(Document(
                name=doc.name,
                text=doc.text,
                annotations=_bio_to_standoff(tags, doc.spacy_doc)
            ))
        except Exception as e:
            logger.warning('Could not convert document to standoff {}\n tags = {}\n{}'
                           .format(doc.name, tags, e))

            annotated_docs.append(Document(
                name=doc.name,
                text=doc.text,
                annotations=[]
            ))

    return annotated_docs


def _group_sentences(sentence_tags: List[List[str]],
                     docs: List[ParsedDoc]) -> List[Tuple[ParsedDoc, List[str]]]:
    """Group BIO tagged sentences by document (i.e., merge list of sentences to a single list per
    document).

    This function is used to convert sentence level BIO annotations to standoff format.

    Example:
    ```
    sentence_tags = [['B', 'I', 'O'], ['O', 'O'], ['O', 'O']]
    docs = [d1, d1, d2]

    _group_sentences(sentence_tags, docs)
    -> [['B', 'I', 'O', 'O', 'O'], ['O', 'O']]
    -> [d1, d2]
    ```

    Parameters
    ----------
    sentence_tags : List[List[str]]
        List of sentences. Each sentence contains a list of BIO tags.
    docs : List[ParsedDoc]
        List of documents corresponding to the sentences.

    Returns
    -------
    grouped_sentences : List[Tuple(ParsedDoc, List[str])]
        The sentences grouped by document.

    """
    grouped = defaultdict(list)
    for doc, sentence in zip(docs, sentence_tags):
        grouped[doc].extend(sentence)

    return list(zip(grouped.keys(), grouped.values()))


def _end_of_chunk(next_tag: str) -> bool:
    return not next_tag or next_tag.startswith('O') or next_tag.startswith('B')


def _bio_to_biluo(bio_tags: List[str]) -> List[str]:
    """Convert sentece-level BIO tagging to BILUO.
    """
    biluo_tags = []
    for i in range(len(bio_tags)):

        current_tag = bio_tags[i]
        if i + 1 < len(bio_tags):
            next_tag = bio_tags[i + 1]
        else:
            next_tag = None

        if current_tag.startswith('B-') and _end_of_chunk(next_tag):
            biluo_tags.append('U-' + current_tag[2:])
        elif current_tag.startswith('I-') and _end_of_chunk(next_tag):
            biluo_tags.append('L-' + current_tag[2:])
        else:
            biluo_tags.append(current_tag)

    return biluo_tags


def _bio_to_standoff(bio_tags: List[str], spacy_doc: spacy.tokens.Doc) -> List[Annotation]:
    """Convert BIO tagged document to annotations in standoff format.

    The original spaCy document is used to recreate correct entity offsets.

    Parameters
    ----------
    bio_tags : List[str]
        A BIO tagged sentence. `len(bio_tags) == len(spacy_doc)` has to hold.
    spacy_doc : spacy.tokens.Doc
        The spaCy doc corresponding to the BIO tags.

    Returns
    -------
    List[Annotation]
        The standoff annotations.

    """
    bio_tags = fix_dangling_entities(bio_tags)
    biluo_tags = _bio_to_biluo(bio_tags)
    offsets = offsets_from_biluo_tags(spacy_doc, biluo_tags)

    annotations = []
    for i, offset in enumerate(offsets):
        annotations.append(Annotation(
            text=spacy_doc.char_span(offset[0], offset[1]).text,
            start=offset[0],
            end=offset[1],
            tag=offset[2],
            ann_id='T{}'.format(i),
        ))

    return annotations


def _doc_to_bio(parsed_doc: spacy.tokens.Doc, annotations: List[Annotation]):
    entities = [(int(ann.start), int(ann.end), ann.tag) for ann in annotations]
    biluo_tags = biluo_tags_from_offsets(parsed_doc, entities)

    biluo_to_bio = {
        'B-': 'B-',
        'I-': 'I-',
        'L-': 'I-',
        'U-': 'B-',
    }

    tags = []
    for tag in biluo_tags:
        if tag == "O":
            tags.append('O')
        elif tag == '-':
            # Returned by spacy if token boundaries mismatch entity boundaries.
            # https://spacy.io/api/goldparse#biluo_tags_from_offsets
            tags.append('O')
            warnings.warn(
                'Some entities could not be aligned in the text. Use `spacy.gold.biluo_tags_from_offsets(nlp.make_doc(text), entities)` to check the alignment.',
                UserWarning
            )
        else:
            tags.append(biluo_to_bio[tag[0:2]] + tag[2:])

    return tags


def _validate_labels(docs: List[ParsedDoc], sents: List[List[Token]]):
    """Perform sanity checks on BIO tagged sentences.

    * A sentence has to start with either a 'B' or 'O' token.
    * If a token is within an entity (i.e., tagged as 'I'), the preceeding token has to be tagged
    as 'I' or 'B'.

    """
    for doc, sent in zip(docs, sents):
        try:
            assert sent[0].label == 'O' or sent[0].label.startswith('B-')
            for i, token in enumerate(sent[1:]):
                if token.label.startswith('I-'):
                    assert sent[i].label.startswith('B-') or sent[i].label.startswith('I-')
        except:
            logger.warning('Invalid tagging {}, sent={}'.format(doc.name, sent))


def _inside_entity(previous_tag, current_tag):
    return previous_tag != 'O' and previous_tag[2:] == current_tag[2:]


def fix_dangling_entities(tags: List[str]) -> List[str]:
    prev_tag = '<OOB>'
    result = []

    for tag in tags:
        if tag.startswith('I-') and not _inside_entity(prev_tag, tag):
            result.append('B-' + tag[2:])
        else:
            result.append(tag)
        prev_tag = tag

    return result
Utility methods to convert between standoff and BIO format. # spacy>=3 # Silence spaCy warning regarding misaligned entity boundaries. It will show up multiple times # because the message changes with the input text. # More info on the warning: https://github.com/explosion/spaCy/issues/5727 Convert corpus into list of BIO tagged sentences.      Each document is parsed using the spaCy tokenizer and segmented into sentences. Afterwards, each     sentence is represented as a list of BIO tagged tokens.      Parameters     ----------     docs : List[Document]         The corpus documents.     tokenizer : Tokenizer         A tokenizer instance that parsed the documents using spaCy.      Returns     -------     sents : List[List[Token]]         List of sentences, where each sentence is a list of BIO tagged tokens.     sents_docs : List[ParsedDoc]         The parsed documents where the sentence originates from. # merge sentences where entities cross sentence boundaries Convert a BIO tagged documents to standoff annotated documents.      Parameters     ----------     sentence_tags : List[List[str]]         List of sentences of BIO tagged tokens.     docs : List[ParsedDoc]         The documents corresponding to each sentence.      Returns     -------     annotated_docs : List[Document]         The documents with annotated entities in standoff format. Group BIO tagged sentences by document (i.e., merge list of sentences to a single list per     document).      This function is used to convert sentence level BIO annotations to standoff format.      Example:     ```     sentence_tags = [['B', 'I', 'O'], ['O', 'O'], ['O', 'O']]     docs = [d1, d1, d2]      _group_sentences(sentence_tags, docs)     -> [['B', 'I', 'O', 'O', 'O'], ['O', 'O']]     -> [d1, d2]     ```      Parameters     ----------     sentence_tags : List[List[str]]         List of sentences. Each sentence contains a list of BIO tags.     docs : List[ParsedDoc]         List of documents corresponding to the sentences.      Returns     -------     grouped_sentences : List[Tuple(ParsedDoc, List[str])]         The sentences grouped by document. Convert sentece-level BIO tagging to BILUO. Convert BIO tagged document to annotations in standoff format.      The original spaCy document is used to recreate correct entity offsets.      Parameters     ----------     bio_tags : List[str]         A BIO tagged sentence. `len(bio_tags) == len(spacy_doc)` has to hold.     spacy_doc : spacy.tokens.Doc         The spaCy doc corresponding to the BIO tags.      Returns     -------     List[Annotation]         The standoff annotations. # Returned by spacy if token boundaries mismatch entity boundaries. # https://spacy.io/api/goldparse#biluo_tags_from_offsets Perform sanity checks on BIO tagged sentences.      * A sentence has to start with either a 'B' or 'O' token.     * If a token is within an entity (i.e., tagged as 'I'), the preceeding token has to be tagged     as 'I' or 'B'.