import gradio as gr
from fastcoref import FCoref
import spacy
from spacy import displacy
import torch
from random import randint
import os

# Run OS scripts to download Ukrainian spaCy
os.system("python -m spacy download uk_core_news_md")

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()