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/IELTS_essay_scoring"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Streamlit app
st.title("Automated Scoring IELTS App")
st.write("Enter your IELTS 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", padding=True, truncation=True, max_length=512)
        
        # 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()
        
        # Apply a significant uniform reduction (e.g., reduce by 80%)
        reduction_factor = 0.6  # Reduce scores by 80%
        adjusted_scores = predicted_scores * reduction_factor
        
        # Ensure scores do not go below zero
        adjusted_scores = np.maximum(adjusted_scores, 0)
        
        # Normalize the scores to ensure they fall within the 0-9 range
        normalized_scores = (adjusted_scores / adjusted_scores.max()) * 9  # Scale to 9
        
        # Apply additional reductions to all scores
        additional_reduction = 1.5  # Further reduce all scores
        normalized_scores = np.maximum(normalized_scores - additional_reduction, 0)
        
        # Round the scores
        rounded_scores = np.round(normalized_scores * 2) / 2
        
        # Display the predictions
        labels = ["Task Achievement", "Coherence and Cohesion", "Vocabulary", "Grammar", "Overall"]
        for label, score in zip(labels, rounded_scores):
            st.write(f"{label}: {score:.1f}")
    else:
        st.write("Please enter some text to get scores.")