import time
import streamlit as st

from flair.data import Sentence
from flair.models import SequenceTagger
from flair.visual.ner_html import render_ner_html

checkpoints = [
    "qanastek/pos-french",
]

@st.cache(suppress_st_warning=True, allow_output_mutation=True)
def get_model(model_name):
    
    # Load the model
    return SequenceTagger.load(model_name)


def main():

    st.title("🥖 French-Part-Of-Speech-Tagging")

    checkpoint = st.selectbox("Choose model", checkpoints)
    model = get_model(checkpoint)

    default_text = "George Washington est allé à Washington"
    input_text = st.text_area(
        label="Original text",
        value=default_text,
    )

    start = None
    if st.button("🧠 Compute"):
        start = time.time()
        with st.spinner("Search for Part-Of-Speech Tags 🔍"):
            
            # Build Sentence
            s = Sentence(input_text)

            # predict tags
            model.predict(s)

            # # print predicted pos tags
            # result = sentence.to_tagged_string()

            try:
                st.write(render_ner_html(sentences=[s], wrap_page=False), unsafe_allow_html=True)
            except Exception as e:
                st.error("Some error occured!" + str(e))
                st.stop()

    st.write("---")
    st.markdown(
        "Built by [Yanis Labrak](https://www.linkedin.com/in/yanis-labrak-8a7412145/) 🚀"
    )
    st.markdown(
        "_Source code made with [FlairNLP](https://github.com/flairNLP/flair)_"
    )

    if start is not None:
        st.text(f"prediction took {time.time() - start:.2f}s")


if __name__ == "__main__":
    main()