regel-corpus/flair-relation-regel2-gene

Info

• Flair NLP model for binary non-directional relation classification: gene<->regulatory element
• pre-trained language model: michiyasunaga/BioLinkBERT-base
• fine-tuned on RegEl2 corpus

Demo: How to use in Flair

Requires flair>=0.15.1

import itertools

from flair.nn import Classifier

from flair.data import Label, Sentence, Span
from flair.datasets import FlairDatapointDataset
from flair.models.relation_classifier_model import (
    EncodedSentence,
    EncodingStrategy,
    RelationClassifier,
)

class NodirEntityMask(EncodingStrategy):
    def __init__(self, entity_types: set[str], *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.special_tokens = {f"[{e}]" for e in entity_types}

    def _encode(self, label: Label) -> str:
        return f"[{label.value}]"

    def encode_head(self, head: Span, label: Label) -> str:
        return self._encode(label)

    def encode_tail(self, tail: Span, label: Label) -> str:
        return self._encode(label)

def get_indices_nondirectional(
    source: list[Sentence] | list[EncodedSentence] | FlairDatapointDataset,
    zero_tag_value: str = "O",
) -> list[int]:
    filtered = {}
    for i, s in enumerate(source):  # type: ignore
        key = s.text
        if key not in filtered:
            filtered[key] = (s, i)
        else:
            # prefer labeled instances over unlabeled
            if filtered[key][0].tag == zero_tag_value:
                filtered[key] = (s, i)

    return [x[1] for x in filtered.values()]


def predict_nondirectional(sentences: list[Sentence], model: RelationClassifier):
    sentences_with_relation_reference = list(
        itertools.chain.from_iterable(
            model._encode_sentence_for_inference(sentence) for sentence in sentences
        )
    )
    encoded_sentences = [x[0] for x in sentences_with_relation_reference]
    indices = get_indices_nondirectional(encoded_sentences)
    encoded_sentences = [encoded_sentences[i] for i in indices]
    model.predict(encoded_sentences)

    # For each encoded sentence, transfer its prediction onto the original relation
    for i, (
        encoded_sentence,
        original_relation,
    ) in enumerate(sentences_with_relation_reference):
        if i in indices:
            for label in encoded_sentence.get_labels(model.label_type):
                original_relation.add_label(
                    model.label_type,
                    value=label.value,
                    score=label.score,
                )

def predict_ner(sentences: list[Sentence], models: dict[str, Classifier]):
    for tag, model in models.items():
        model.predict(sentences, label_name=tag)

    for sentence in sentences:
        for tag in models.keys():
            spans = sentence.get_spans(label_type=tag)
            for span in spans:
                start = span.tokens[0].idx - 1
                end = span.tokens[-1].idx
                sentence[start:end].add_label("ner", tag)
            sentence.remove_labels(tag)

def predict_nondirectional(sentences: list[Sentence], model: RelationClassifier):
    sentences_with_relation_reference = list(
        itertools.chain.from_iterable(
            model._encode_sentence_for_inference(sentence) for sentence in sentences
        )
    )
    encoded_sentences = [x[0] for x in sentences_with_relation_reference]
    indices = get_indices_nondirectional(encoded_sentences)
    encoded_sentences = [encoded_sentences[i] for i in indices]
    model.predict(encoded_sentences)

    # For each encoded sentence, transfer its prediction onto the original relation
    for i, (
        encoded_sentence,
        original_relation,
    ) in enumerate(sentences_with_relation_reference):
        if i in indices:
            for label in encoded_sentence.get_labels(model.label_type):
                original_relation.add_label(
                    model.label_type,
                    value=label.value,
                    score=label.score,
                )


text = "Transient transfection analysis demonstrated that PU.1 functions to repress the IgH intronic enhancer"
sentences = [Sentence(text, use_tokenizer=SciSpacyTokenizer())]


ner_models = {
    "ENHANCER": Classifier.load('regel-corpus/hunflair2-regel2-enhancer'),
    "PROMOTER": Classifier.load('regel-corpus/hunflair2-regel2-promoter'),
    "TFBS": Classifier.load('regel-corpus/hunflair2-regel2-tfbs'),
}
predict_ner(sentences=sentences, models=ner_models)


rc_model = RelationClassifier.load('regel-corpus/flair-relation-regel2-gene')

# to avoid issues in loading, the model was saved with the encoding strategy `TypedEntityMarker` 
encoding_strategy = NodirEntityMask(entity_types=rc_model.entity_label_types["ner"])  # type: ignore
rc_model.encoding_strategy = encoding_strategy

# prediction logic in flair is directional, i.e. it distinguishes between (A, relation, B) and (B, relation, A)
predict_nondirectional(sentences=sentences, model=rc_model)