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):
    doc = nlp(text)
    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.pos_}</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 dan, danas."]
)

iface.launch()