|
|
|
import os |
|
|
|
os.system("pip3 install torch==1.10.1+cpu torchvision==0.11.2+cpu torchaudio==0.10.1+cpu -f " |
|
"https://download.pytorch.org/whl/cpu/torch_stable.html") |
|
|
|
import gradio as gr |
|
from transformers import pipeline |
|
|
|
import spacy |
|
from spacy import displacy |
|
|
|
ner_map = {0: '0', |
|
1: 'B-OSOBA', |
|
2: 'I-OSOBA', |
|
3: 'B-ORGANIZÁCIA', |
|
4: 'I-ORGANIZÁCIA', |
|
5: 'B-LOKALITA', |
|
6: 'I-LOKALITA'} |
|
|
|
options = {"ents": ["OSOBA", |
|
"ORGANIZÁCIA", |
|
"LOKALITA"], |
|
"colors": {"OSOBA": "lightblue", |
|
"ORGANIZÁCIA": "lightcoral", |
|
"LOKALITA": "lightgreen"}} |
|
|
|
ner_pipeline = pipeline(task='ner', model="crabz/slovakbert-ner") |
|
nlp = spacy.blank("sk") |
|
|
|
|
|
def apply_ner(sentence: str): |
|
classifications = ner_pipeline(sentence) |
|
|
|
entities = [] |
|
for i in range(len(classifications)): |
|
if classifications[i]['entity'] != 0: |
|
if ner_map[classifications[i]['entity']][0] == 'B': |
|
j = i + 1 |
|
while j < len(classifications) and ner_map[classifications[j]['entity']][0] == 'I': |
|
j += 1 |
|
entities.append((ner_map[classifications[i]['entity']].split('-')[1], classifications[i]['start'], |
|
classifications[j - 1]['end'])) |
|
doc = nlp(sentence) |
|
|
|
ents = [] |
|
for ee in entities: |
|
ents.append(doc.char_span(ee[1], ee[2], ee[0])) |
|
doc.ents = ents |
|
|
|
displacy_html = displacy.render(doc, style="ent", options=options) |
|
return displacy_html |
|
|
|
|
|
intf = gr.Interface(fn=apply_ner, inputs="text", outputs="html", title='Slovak Named Entity Recognition', |
|
allow_flagging=False, |
|
examples=[["Laboratóriá Úradu verejného zdravotníctva sekvenovaním potvrdili výskyt ďalších " |
|
"štyroch prípadov variantu omikron na Slovensku."], |
|
["Čaputová opakovane tvrdí, že \"spravodlivosť na Slovensku neplatí vždy pre všetkých " |
|
"rovnako\"."]], |
|
description="Named Entity Recognition (NER) labels persons, organizations and locations in the " |
|
"given sentence. You can try out one of the examples below or type your own sentence. " |
|
"Don't forget to use double quotes (\") instead of curved quotes („“)", |
|
article="") |
|
intf.launch() |
|
|
