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)
def load_model(input_text):
    tokenizer = AutoTokenizer.from_pretrained(cfg["model_name_or_path"])
    model = RobertaForSequenceClassification.from_pretrained(cfg["model_name_or_path"])

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


st.title("RoBERTa Marathi")

input_text = st.text_input("Text:")

predict_button = st.button("Predict")

if predict_button:
    with st.spinner("Generating prediction..."):
        # Get prediction here
        result = load_model(input_text)
        st.write(result)