import gradio as gr
import spacy
from spacy import displacy
from spacy.tokens import Span
from random import randint
from fastcoref import FCoref

model = FCoref("pythainlp/han-coref-v1.0")
nlp = spacy.blank("th")

default = "แกว่าเราเข้มแข็งมากพอที่จะเห็นแกไปน่ารักกับใครหรอ"

def corefer(text):
    preds = model.predict(texts=[text])
    clusters = preds[0].get_clusters(as_strings=False)
    title=None
    dic_ents = {"text":text,"ents":[],"title":title}
    _tag=[str(i) for i in list(range(len(clusters)))]
    for i,tag in enumerate(clusters):
        for s,e in tag:
            dic_ents["ents"].append({"start": s, "end": e, "label": _tag[i]})
    colors={i:"#"+''.join([random.choice('0123456789ABCDEF') for j in range(6)]) for i in _tag} # thank https://stackoverflow.com/a/50218895
    return displacy.render(dic_ents, manual=True, style="ent",options=colors)


iface = gr.Interface(fn=corefer, 
                     inputs=gr.Textbox(label="Enter Text To Corefer with FastCoref", lines=2, value=default), 
                     outputs="html").queue()
iface.launch()