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)
    doc = nlp(text)
    doc.spans["sc"] = []
    colors = {"Cluster {}".format(i):'#%06X' % randint(0, 0xFFFFFF) 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 with FastCoref", lines=2, value=default), 
                     outputs="html")
iface.launch()