import json

import streamlit as st
from transformers import AutoTokenizer, RobertaForSequenceClassification, pipeline

with open("config.json") as f:
    cfg = json.loads(f.read())


@st.cache(allow_output_mutation=True, show_spinner=False)
def load_model(input_text, model_name_or_path):
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
    model = RobertaForSequenceClassification.from_pretrained(model_name_or_path)

    nlp = pipeline("text-classification", model=model, tokenizer=tokenizer)
    result = nlp(input_text)
    return result


def app():
    st.title("RoBERTa Marathi")

    st.markdown(
        "This demo uses [RoBERTa for Marathi](https://huggingface.co/flax-community/roberta-base-mr) model "
        "trained on [mC4](https://huggingface.co/datasets/mc4)."
    )

    st.markdown(
        "Can't figure out where to get a sample text? Visit any of the following links, copy any headline and see if "
        "the model is predicting the respective class or not.\n"
        "- [entertainment](https://maharashtratimes.com/entertainment/articlelist/19359255.cms)\n"
        "- [sports](https://maharashtratimes.com/sports/articlelist/2429056.cms)\n"
        "- [lifestyle](https://maharashtratimes.com/lifestyle-news/articlelist/2429025.cms)"
        
    )

    classifier = st.sidebar.selectbox("Select a Model", index=0, options=["Indic NLP", "iNLTK"])

    st.sidebar.markdown(
        "- [Indic NLP](https://huggingface.co/flax-community/mr-indicnlp-classifier) fine-tuned on "
        "[Indic NLP Corpus for Marathi](https://github.com/ai4bharat/indicnlp_corpus#indicnlp-news-article-classification-dataset)"
        "\n"
        "- [iNLTK](https://huggingface.co/flax-community/mr-inltk-classifier) fine-tuned on "
        "[Marathi News Dataset](https://www.kaggle.com/disisbig/marathi-news-dataset)"
    )

    sample_texts = [
        "अध्यक्ष शरद पवार आणि उपमुख्यमंत्री अजित पवार यांची भेट घेतली.",
        "मोठी बातमी! उद्या दुपारी १ वाजता जाहीर होणार दहावीचा निकाल",
        "Custom",
    ]
    model_name_or_path = cfg["models"][classifier]

    text_to_classify = st.selectbox("Select a Text", options=sample_texts, index=len(sample_texts) - 1)

    if text_to_classify == "Custom":
        text_to_classify = st.text_input("Enter custom text:")

    predict_button = st.button("Predict")

    if predict_button:
        with st.spinner("Generating prediction..."):
            result = load_model(text_to_classify, model_name_or_path)

            st.markdown("**Predicted label:** " + result[0]["label"])
            st.markdown("**Score:** " + str(round(result[0]["score"], 2)))