import spacy
import gradio as gr
from spacy import displacy

# Load your trained spaCy model
nlp = spacy.load("sr_Spacy_Serbian_Model_SrpKor4Tagging_BERTICOVO")

# Define a function to display the tags and lemmas
def display_tags_and_lemmas(text):
    # First, we'll use spaCy to tag and parse the text
    text = text.strip()
    if not text:
        return "", ""
    doc = nlp(text.strip())
    html = displacy.render(doc, style="ent", page=True, minify=True)

    # We'll also create a custom HTML to display lemmas nicely
    lemma_html = "<div class='lemma-display'><table class='table table-hover'>"
    lemma_html += "<thead><tr><th>Token</th><th>Lemma</th><th>POS Tag</th></tr></thead><tbody>"

    for token in doc:
        lemma_html += f"<tr><td>{token.text}</td><td>{token.lemma_}</td><td>{token.tag_}</td></tr>"

    lemma_html += "</tbody></table></div>"

    # Return both the displaCy HTML and our custom lemma table
    return html, lemma_html

# Define Gradio interface
iface = gr.Interface(
    fn=display_tags_and_lemmas,
    inputs=gr.Textbox(lines=5, placeholder="Unesite rečenicu ovde..."),
    outputs=gr.HTML(label="Leme i POS oznake"),
    title="spaCy Tagger i Lemmatizer",
    description="Unesite rečenicu da biste videli njene imenovane entitete, POS oznake i leme.",
    examples=["Lep je dan, danas. Sutra će biti još lepši!", "Psi su trčali svakog dana. Mačke su spavale."
             "Sedam dana nije dugo."],
    theme="compact",)

iface.launch()