f-coref-ua / app.py
artem-oneai
Add app files
27e9cd8
raw
history blame
1.33 kB
import gradio as gr
from fastcoref import FCoref
import spacy
from spacy import displacy
import torch
from random import randint
nlp = spacy.load('uk_core_news_md')
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
model_path = "artemkramov/coref-ua"
model = FCoref(model_name_or_path=model_path, device=device, nlp=nlp)
default = "Мій друг дав мені свою машину та ключі до неї; крім того, він дав мені його книгу. Я з радістю її читаю."
def random_color():
rand = lambda: randint(100, 255)
return '#%02X%02X%02X' % (rand(), rand(), rand())
def corefer(text):
preds = model.predict(texts=[text])
clusters = preds[0].get_clusters(as_strings=False)
doc = nlp(text)
doc.spans["sc"] = []
colors = {"Cluster {}".format(i): random_color() for i in range(len(clusters))}
for i, cluster in enumerate(clusters):
for sp in cluster:
doc.spans["sc"] += [doc.char_span(sp[0], sp[1], "Cluster {}".format(i))]
return displacy.render(doc, style="span", options={"colors": colors}, page=True)
iface = gr.Interface(fn=corefer,
inputs=gr.Textbox(label="Enter Text To Corefer", lines=2, value=default),
outputs="html", allow_flagging="never")
iface.launch()