import streamlit as st
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
import numpy as np

# Load the model and tokenizer from Hugging Face
model_name = "KevSun/Engessay_grading_ML"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Streamlit app
st.title("Automated Scoring App")
st.write("Enter your English essay below to predict scores from multiple dimensions:")

# Input text from user
user_input = st.text_area("Your text here:")

if st.button("Predict"):
    if user_input:
        # Tokenize input text
        inputs = tokenizer(user_input, return_tensors="pt")
        
        # Get predictions from the model
        with torch.no_grad():
            outputs = model(**inputs)
        
        # Extract the predictions
        predictions = outputs.logits.squeeze()

        # Convert to numpy array if necessary
        predicted_scores = predictions.numpy()
        #predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
        #predictions = predictions[0].tolist()
        
        # Convert predictions to a NumPy array for the calculations
        #predictions_np = np.array(predictions)
        
        # Scale the predictions
        scaled_scores = 2.25 * predicted_scores - 1.25
        rounded_scores = [round(score * 2) / 2 for score in scaled_scores]  # Round to nearest 0.5

        # Display the predictions
        labels = ["cohesion", "syntax", "vocabulary", "phraseology", "grammar", "conventions"]
        for label, score in zip(labels, rounded_scores):
            st.write(f"{label}: {score:}")
    else:
        st.write("Please enter some text to get scores.")