def _get_tars_formatted_sentence(self, label, sentence):
original_text = sentence.to_tokenized_string()
label_text_pair = (f"{label} {self.separator} {original_text}"
if self.prefix else
f"{original_text} {self.separator} {label}")
label_length = 0 if not self.prefix else len(label.split(" ")) + len(
self.separator.split(" "))
# make a tars sentence where all labels are O by default
tars_sentence = Sentence(label_text_pair, use_tokenizer=False)
for entity_label in sentence.get_labels(self.label_type):
if entity_label.value == label:
new_span = [
tars_sentence.get_token(token.idx + label_length)
for token in entity_label.span
]
tars_sentence.add_complex_label(
self.static_label_type,
SpanLabel(Span(new_span), value="entity"))
return tars_sentence
def token_list_to_sentence(self, token_list: conllu.TokenList) -> Sentence:
sentence: Sentence = Sentence()
# current token ID
token_idx = 0
for conllu_token in token_list:
token = Token(conllu_token["form"])
if "ner" in conllu_token:
token.add_label("ner", conllu_token["ner"])
if "ner-2" in conllu_token:
token.add_label("ner-2", conllu_token["ner-2"])
if "lemma" in conllu_token:
token.add_label("lemma", conllu_token["lemma"])
if "misc" in conllu_token and conllu_token["misc"] is not None:
space_after = conllu_token["misc"].get("SpaceAfter")
if space_after == "No":
token.whitespace_after = False
sentence.add_token(token)
token_idx += 1
if "sentence_id" in token_list.metadata:
sentence.add_label("sentence_id", token_list.metadata["sentence_id"])
if "relations" in token_list.metadata:
# relations: List[Relation] = []
for head_start, head_end, tail_start, tail_end, label in token_list.metadata["relations"]:
# head and tail span indices are 1-indexed and end index is inclusive
head = Span(sentence.tokens[head_start - 1 : head_end])
tail = Span(sentence.tokens[tail_start - 1 : tail_end])
sentence.add_complex_label("relation", RelationLabel(value=label, head=head, tail=tail))
# determine all NER label types in sentence and add all NER spans as sentence-level labels
ner_label_types = []
for token in sentence.tokens:
for annotation in token.annotation_layers.keys():
if annotation.startswith("ner") and annotation not in ner_label_types:
ner_label_types.append(annotation)
for label_type in ner_label_types:
spans = sentence.get_spans(label_type)
for span in spans:
sentence.add_complex_label("entity", label=SpanLabel(span=span, value=span.tag, score=span.score))
return sentence
def _label(self, sentence: Sentence):
"""
This will add a complex_label to the given sentence for every match.span() for every registered_mapping.
If a match span overlaps with a token span an exception is raised.
"""
collection = RegexpTagger.TokenCollection(sentence)
for label, pattern in self._regexp_mapping.items():
for match in pattern.finditer(sentence.to_original_text()):
span: Tuple[int, int] = match.span()
try:
token_span = collection.get_token_span(span)
except ValueError:
raise Exception(
f"The match span {span} for label '{label}' is overlapping with a token!"
)
sentence.add_complex_label(label, SpanLabel(token_span, label))
def forward_pass(
self,
sentences: Union[List[Sentence], Sentence],
return_label_candidates: bool = False,
):
if not isinstance(sentences, list):
sentences = [sentences]
# filter sentences with no candidates (no candidates means nothing can be linked anyway)
filtered_sentences = []
for sentence in sentences:
if sentence.get_labels(self.label_type):
filtered_sentences.append(sentence)
# fields to return
span_labels = []
sentences_to_spans = []
empty_label_candidates = []
# if the entire batch has no sentence with candidates, return empty
if len(filtered_sentences) == 0:
scores = None
# otherwise, embed sentence and send through prediction head
else:
# embed all tokens
self.word_embeddings.embed(filtered_sentences)
embedding_names = self.word_embeddings.get_names()
embedding_list = []
# get the embeddings of the entity mentions
for sentence in filtered_sentences:
entities = sentence.get_labels(self.label_type)
for entity in entities:
if self.skip_unk_probability and self.training and entity.value not in self.known_entities:
sample = random.uniform(0, 1)
if sample < self.skip_unk_probability:
continue
span_labels.append([entity.value])
if self.pooling_operation == "first&last":
mention_emb = torch.cat(
(
entity.span.tokens[0].get_embedding(
embedding_names),
entity.span.tokens[-1].get_embedding(
embedding_names),
),
0,
)
embedding_list.append(mention_emb.unsqueeze(0))
if return_label_candidates:
sentences_to_spans.append(sentence)
candidate = SpanLabel(span=entity.span,
value=None,
score=0.0)
empty_label_candidates.append(candidate)
if len(embedding_list) > 0:
embedding_tensor = torch.cat(embedding_list,
0).to(flair.device)
if self.use_dropout:
embedding_tensor = self.dropout(embedding_tensor)
scores = self.decoder(embedding_tensor)
else:
scores = None
if return_label_candidates:
return scores, span_labels, sentences_to_spans, empty_label_candidates
return scores, span_labels
def forward_pass(
self,
sentences: Union[List[Sentence], Sentence],
return_label_candidates: bool = False,
):
if not isinstance(sentences, list):
sentences = [sentences]
# filter sentences with no candidates (no candidates means nothing can be linked anyway)
filtered_sentences = []
for sentence in sentences:
if sentence.get_labels(self.label_type):
filtered_sentences.append(sentence)
# fields to return
span_labels = []
sentences_to_spans = []
empty_label_candidates = []
# if the entire batch has no sentence with candidates, return empty
if len(filtered_sentences) == 0:
scores = None
# otherwise, embed sentence and send through prediction head
else:
# embed all tokens
self.word_embeddings.embed(filtered_sentences)
embedding_names = self.word_embeddings.get_names()
embedding_list = []
# get the embeddings of the entity mentions
for sentence in filtered_sentences:
spans = sentence.get_spans(self.label_type)
for span in spans:
mention_emb = torch.Tensor(
0,
self.word_embeddings.embedding_length).to(flair.device)
for token in span.tokens:
mention_emb = torch.cat(
(
mention_emb,
token.get_embedding(embedding_names).unsqueeze(
0),
),
0,
)
embedding_list.append(
self.aggregated_embedding(mention_emb).unsqueeze(0))
span_labels.append([
label.value
for label in span.get_labels(typename=self.label_type)
])
if return_label_candidates:
sentences_to_spans.append(sentence)
candidate = SpanLabel(span=span, value=None, score=0.0)
empty_label_candidates.append(candidate)
embedding_tensor = torch.cat(embedding_list, 0).to(flair.device)
scores = self.decoder(embedding_tensor)
if return_label_candidates:
return scores, span_labels, sentences_to_spans, empty_label_candidates
return scores, span_labels
def predict(
self,
sentences: Union[List[Sentence], Sentence],
mini_batch_size: int = 32,
return_probabilities_for_all_classes: bool = False,
verbose: bool = False,
label_name: Optional[str] = None,
return_loss=False,
embedding_storage_mode="none",
):
"""
Predicts labels for current batch with CRF or Softmax.
:param sentences: List of sentences in batch
:param mini_batch_size: batch size for test data
:param return_probabilities_for_all_classes: Whether to return probabilites for all classes
:param verbose: whether to use progress bar
:param label_name: which label to predict
:param return_loss: whether to return loss value
:param embedding_storage_mode: determines where to store embeddings - can be "gpu", "cpu" or None.
"""
if label_name is None:
label_name = self.tag_type
with torch.no_grad():
if not sentences:
return sentences
# make sure its a list
if not isinstance(sentences, list) and not isinstance(
sentences, flair.data.Dataset):
sentences = [sentences]
# filter empty sentences
sentences = [
sentence for sentence in sentences if len(sentence) > 0
]
# reverse sort all sequences by their length
reordered_sentences = sorted(sentences,
key=lambda s: len(s),
reverse=True)
if len(reordered_sentences) == 0:
return sentences
dataloader = DataLoader(
dataset=FlairDatapointDataset(reordered_sentences),
batch_size=mini_batch_size,
)
# progress bar for verbosity
if verbose:
dataloader = tqdm(dataloader)
overall_loss = 0
batch_no = 0
label_count = 0
for batch in dataloader:
batch_no += 1
if verbose:
dataloader.set_description(
f"Inferencing on batch {batch_no}")
# stop if all sentences are empty
if not batch:
continue
# get features from forward propagation
features, gold_labels = self.forward(batch)
# remove previously predicted labels of this type
for sentence in batch:
sentence.remove_labels(label_name)
# if return_loss, get loss value
if return_loss:
loss = self._calculate_loss(features, gold_labels)
overall_loss += loss[0]
label_count += loss[1]
# Sort batch in same way as forward propagation
lengths = torch.LongTensor(
[len(sentence) for sentence in batch])
lengths = lengths.sort(dim=0, descending=True)
batch = [batch[i] for i in lengths.indices]
# make predictions
if self.use_crf:
predictions, all_tags = self.viterbi_decoder.decode(
features, return_probabilities_for_all_classes)
else:
predictions, all_tags = self._standard_inference(
features, batch, return_probabilities_for_all_classes)
# add predictions to Sentence
for sentence, sentence_predictions in zip(batch, predictions):
# BIOES-labels need to be converted to spans
if self.predict_spans:
sentence_tags = [
label.value for label in sentence_predictions
]
sentence_scores = [
label.score for label in sentence_predictions
]
predicted_spans = get_spans_from_bio(
sentence_tags, sentence_scores)
for predicted_span in predicted_spans:
span = Span(sentence[
predicted_span[0][0]:predicted_span[0][-1] +
1])
sentence.add_complex_label(
typename=label_name,
label=SpanLabel(span=span,
value=predicted_span[2],
score=predicted_span[1]),
)
# token-labels can be added directly
else:
for token, label in zip(sentence.tokens,
sentence_predictions):
token.add_tag_label(label_name, label)
# all_tags will be empty if all_tag_prob is set to False, so the for loop will be avoided
for (sentence, sent_all_tags) in zip(batch, all_tags):
for (token, token_all_tags) in zip(sentence.tokens,
sent_all_tags):
token.add_tags_proba_dist(label_name, token_all_tags)
store_embeddings(sentences, storage_mode=embedding_storage_mode)
if return_loss:
return overall_loss, label_count
def token_list_to_sentence(self, token_list: conllu.TokenList) -> Sentence:
sentence: Sentence = Sentence()
# Build the sentence tokens and add the annotations.
for conllu_token in token_list:
token = Token(conllu_token["form"])
for field in self.token_annotation_fields:
field_value: Any = conllu_token[field]
if isinstance(field_value, dict):
# For fields that contain key-value annotations,
# we add the key as label type-name and the value as the label value.
for key, value in field_value.items():
token.add_label(typename=key, value=str(value))
else:
token.add_label(typename=field, value=str(field_value))
if conllu_token.get("misc") is not None:
space_after: Optional[str] = conllu_token["misc"].get(
"SpaceAfter")
if space_after == "No":
token.whitespace_after = False
sentence.add_token(token)
if "sentence_id" in token_list.metadata:
sentence.add_label("sentence_id",
token_list.metadata["sentence_id"])
if "relations" in token_list.metadata:
for (
head_start,
head_end,
tail_start,
tail_end,
label,
) in token_list.metadata["relations"]:
# head and tail span indices are 1-indexed and end index is inclusive
head = Span(sentence.tokens[head_start - 1:head_end])
tail = Span(sentence.tokens[tail_start - 1:tail_end])
sentence.add_complex_label(
"relation", RelationLabel(value=label,
head=head,
tail=tail))
# determine all NER label types in sentence and add all NER spans as sentence-level labels
ner_label_types = []
for token in sentence.tokens:
for annotation in token.annotation_layers.keys():
if annotation.startswith(
"ner") and annotation not in ner_label_types:
ner_label_types.append(annotation)
for label_type in ner_label_types:
spans = sentence.get_spans(label_type)
for span in spans:
sentence.add_complex_label(
"entity",
label=SpanLabel(span=span,
value=span.tag,
score=span.score),
)
return sentence
def predict(
self,
sentences: Union[List[Sentence], Sentence],
mini_batch_size=32,
return_probabilities_for_all_classes: bool = False,
verbose: bool = False,
label_name: Optional[str] = None,
return_loss=False,
embedding_storage_mode="none",
most_probable_first: bool = True,
):
# return
"""
Predict sequence tags for Named Entity Recognition task
:param sentences: a Sentence or a List of Sentence
:param mini_batch_size: size of the minibatch, usually bigger is more rapid but consume more memory,
up to a point when it has no more effect.
:param all_tag_prob: True to compute the score for each tag on each token,
otherwise only the score of the best tag is returned
:param verbose: set to True to display a progress bar
:param return_loss: set to True to return loss
:param label_name: set this to change the name of the label type that is predicted
:param embedding_storage_mode: default is 'none' which is always best. Only set to 'cpu' or 'gpu' if
you wish to not only predict, but also keep the generated embeddings in CPU or GPU memory respectively.
'gpu' to store embeddings in GPU memory.
"""
if label_name is None:
label_name = self.get_current_label_type()
# with torch.no_grad():
if not sentences:
return sentences
if not isinstance(sentences, list):
sentences = [sentences]
reordered_sentences = sorted(sentences,
key=lambda s: len(s),
reverse=True)
dataloader = DataLoader(
dataset=FlairDatapointDataset(reordered_sentences),
batch_size=mini_batch_size,
)
# progress bar for verbosity
if verbose:
dataloader = tqdm(dataloader)
overall_loss = 0
overall_count = 0
with torch.no_grad():
for batch in dataloader:
batch = self._filter_empty_sentences(batch)
# stop if all sentences are empty
if not batch:
continue
# go through each sentence in the batch
for sentence in batch:
# always remove tags first
sentence.remove_labels(label_name)
all_labels = [
label.decode("utf-8") for label in
self.get_current_label_dictionary().idx2item
]
all_detected = {}
for label in all_labels:
tars_sentence = self._get_tars_formatted_sentence(
label, sentence)
loss_and_count = self.tars_model.predict(
tars_sentence,
label_name=label_name,
return_loss=True,
)
overall_loss += loss_and_count[0].item()
overall_count += loss_and_count[1]
for predicted in tars_sentence.get_labels(label_name):
predicted.value = label
all_detected[predicted] = predicted.score
if most_probable_first:
import operator
already_set_indices = []
sorted_x = sorted(all_detected.items(),
key=operator.itemgetter(1))
sorted_x.reverse()
for tuple in sorted_x:
# get the span and its label
label = tuple[0]
# label = span.get_labels("tars_temp_label")[0].value
label_length = (0 if not self.prefix else
len(label.value.split(" ")) +
len(self.separator.split(" ")))
# determine whether tokens in this span already have a label
tag_this = True
for token in label.span:
corresponding_token = sentence.get_token(
token.idx - label_length)
if corresponding_token is None:
tag_this = False
continue
if token.idx in already_set_indices:
tag_this = False
continue
# only add if all tokens have no label
if tag_this:
already_set_indices.extend(
token.idx for token in label.span)
predicted_span = [
sentence.get_token(token.idx -
label_length)
for token in label.span
]
sentence.add_complex_label(
label_name,
label=SpanLabel(Span(predicted_span),
value=label.value,
score=label.score),
)
# clearing token embeddings to save memory
store_embeddings(batch, storage_mode=embedding_storage_mode)
if return_loss:
return overall_loss, overall_count