Views
No views yet
gene<->regulatory element1import itertools
2
3from flair.nn import Classifier
4
5from flair.data import Label, Sentence, Span
6from flair.datasets import FlairDatapointDataset
7from flair.models.relation_classifier_model import (
8 EncodedSentence,
9 EncodingStrategy,
10 RelationClassifier,
11)
12
13class NodirEntityMask(EncodingStrategy):
14 def __init__(self, entity_types: set[str], *args, **kwargs):
15 super().__init__(*args, **kwargs)
16 self.special_tokens = {f"[{e}]" for e in entity_types}
17
18 def _encode(self, label: Label) -> str:
19 return f"[{label.value}]"
20
21 def encode_head(self, head: Span, label: Label) -> str:
22 return self._encode(label)
23
24 def encode_tail(self, tail: Span, label: Label) -> str:
25 return self._encode(label)
26
27def get_indices_nondirectional(
28 source: list[Sentence] | list[EncodedSentence] | FlairDatapointDataset,
29 zero_tag_value: str = "O",
30) -> list[int]:
31 filtered = {}
32 for i, s in enumerate(source): # type: ignore
33 key = s.text
34 if key not in filtered:
35 filtered[key] = (s, i)
36 else:
37 # prefer labeled instances over unlabeled
38 if filtered[key][0].tag == zero_tag_value:
39 filtered[key] = (s, i)
40
41 return [x[1] for x in filtered.values()]
42
43
44def predict_nondirectional(sentences: list[Sentence], model: RelationClassifier):
45 sentences_with_relation_reference = list(
46 itertools.chain.from_iterable(
47 model._encode_sentence_for_inference(sentence) for sentence in sentences
48 )
49 )
50 encoded_sentences = [x[0] for x in sentences_with_relation_reference]
51 indices = get_indices_nondirectional(encoded_sentences)
52 encoded_sentences = [encoded_sentences[i] for i in indices]
53 model.predict(encoded_sentences)
54
55 # For each encoded sentence, transfer its prediction onto the original relation
56 for i, (
57 encoded_sentence,
58 original_relation,
59 ) in enumerate(sentences_with_relation_reference):
60 if i in indices:
61 for label in encoded_sentence.get_labels(model.label_type):
62 original_relation.add_label(
63 model.label_type,
64 value=label.value,
65 score=label.score,
66 )
67
68def predict_ner(sentences: list[Sentence], models: dict[str, Classifier]):
69 for tag, model in models.items():
70 model.predict(sentences, label_name=tag)
71
72 for sentence in sentences:
73 for tag in models.keys():
74 spans = sentence.get_spans(label_type=tag)
75 for span in spans:
76 start = span.tokens[0].idx - 1
77 end = span.tokens[-1].idx
78 sentence[start:end].add_label("ner", tag)
79 sentence.remove_labels(tag)
80
81def predict_nondirectional(sentences: list[Sentence], model: RelationClassifier):
82 sentences_with_relation_reference = list(
83 itertools.chain.from_iterable(
84 model._encode_sentence_for_inference(sentence) for sentence in sentences
85 )
86 )
87 encoded_sentences = [x[0] for x in sentences_with_relation_reference]
88 indices = get_indices_nondirectional(encoded_sentences)
89 encoded_sentences = [encoded_sentences[i] for i in indices]
90 model.predict(encoded_sentences)
91
92 # For each encoded sentence, transfer its prediction onto the original relation
93 for i, (
94 encoded_sentence,
95 original_relation,
96 ) in enumerate(sentences_with_relation_reference):
97 if i in indices:
98 for label in encoded_sentence.get_labels(model.label_type):
99 original_relation.add_label(
100 model.label_type,
101 value=label.value,
102 score=label.score,
103 )
104
105
106text = "Transient transfection analysis demonstrated that PU.1 functions to repress the IgH intronic enhancer"
107sentences = [Sentence(text, use_tokenizer=SciSpacyTokenizer())]
108
109
110ner_models = {
111 "ENHANCER": Classifier.load('regel-corpus/hunflair2-regel2-enhancer'),
112 "PROMOTER": Classifier.load('regel-corpus/hunflair2-regel2-promoter'),
113 "TFBS": Classifier.load('regel-corpus/hunflair2-regel2-tfbs'),
114}
115predict_ner(sentences=sentences, models=ner_models)
116
117
118rc_model = RelationClassifier.load('regel-corpus/flair-relation-regel2-gene')
119
120# to avoid issues in loading, the model was saved with the encoding strategy `TypedEntityMarker`
121encoding_strategy = NodirEntityMask(entity_types=rc_model.entity_label_types["ner"]) # type: ignore
122rc_model.encoding_strategy = encoding_strategy
123
124# prediction logic in flair is directional, i.e. it distinguishes between (A, relation, B) and (B, relation, A)
125predict_nondirectional(sentences=sentences, model=rc_model)