from flair.models import SequenceTagger
from flair.data import Sentence


def predict_entity_from_text(sentence) -> str:
    sentence_tagged = Sentence(sentence)
    tagger.predict(sentence_tagged)
    new_sent = ""
    idx = 0
    for entity in sentence_tagged.get_spans('ner'):
        new_sent += sentence[idx:entity.start_pos]
        replacement = "|" + entity.text + "| (" + entity.tag + ", " + str(round(entity.score, 2)) + ")"
        new_sent += replacement
        idx = entity.end_pos

    new_sent += sentence[idx:]
    return new_sent


tagger = SequenceTagger.load('./best-model.pt')

"""
text = 'George Washington ging nach Washington.'
print(predict_entity_from_text(text))
"""
