import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F

# Load the tokenizer and model from HuggingFace Hub
@st.cache_resource(show_spinner=False)
def load_model():
    model_name = "pavlyhalim/BERT_ALL_README"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    model.eval()
    return tokenizer, model

tokenizer, model = load_model()

st.title("Readability Score Predictor based on BERT")

st.write("""
Enter a sentence, and the model will predict its readability score (from 1 to 6).
""")

user_input = st.text_area("Enter your sentence here:", height=100)

if st.button("Predict Readability Score"):
    if user_input.strip() == "":
        st.warning("Please enter a sentence.")
    else:
        with st.spinner('Predicting...'):
            inputs = tokenizer(
                user_input,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=128
            )

            with torch.no_grad():
                outputs = model(**inputs)
                logits = outputs.logits

            probabilities = F.softmax(logits, dim=1)
            predicted_class = torch.argmax(probabilities, dim=1).item()
            predicted_probability = probabilities[0][predicted_class].item()

            predicted_label = predicted_class + 1

            st.success(f"Predicted Readability Score: **{predicted_label}**")
            st.write(f"Confidence: **{predicted_probability * 100:.2f}%**")

            st.write("### Class Probabilities:")
            for i, prob in enumerate(probabilities[0]):
                label = i + 1  
                st.write(f"Score {label}: {prob.item() * 100:.2f}%")