from typing import *
import gradio as gr
import numpy as np
import spaces
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_name = "JacobLinCool/IELTS_essay_scoring_safetensors"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)


@spaces.GPU
def grade(question: str, answer: str) -> Tuple[float, float, float, float, float]:
    if len(question) < 30 or len(answer) < 30:
        raise gr.Error("Please enter more than 30 characters")

    text = f"{question} {answer}"

    inputs = tokenizer(
        text, return_tensors="pt", padding=True, truncation=True, max_length=512
    )
    with torch.no_grad():
        outputs = model(**inputs)
        predictions = outputs.logits.squeeze()

    predicted_scores = predictions.numpy()
    normalized_scores = (predicted_scores / predicted_scores.max()) * 9
    rounded_scores = np.round(normalized_scores * 2) / 2

    return tuple(rounded_scores)


with gr.Blocks() as app:
    gr.Markdown("# Essay Scoring")

    with gr.Row():
        with gr.Column():
            question = gr.Textbox(
                label="Enter the question here",
                placeholder="Write the question here",
                lines=3,
            )
            essay = gr.Textbox(
                label="Enter your essay here",
                placeholder="Write your essay here",
                lines=10,
            )
            btn = gr.Button("Grade Essay", variant="primary")
        with gr.Column():
            task_achievement = gr.Number(label="Task Achievement")
            coherence_cohesion = gr.Number(label="Coherence and Cohesion")
            vocabulary = gr.Number(label="Vocabulary")
            grammar = gr.Number(label="Grammar")
            overall = gr.Number(label="Overall")

    btn.click(
        fn=grade,
        inputs=[question, essay],
        outputs=[task_achievement, coherence_cohesion, vocabulary, grammar, overall],
    )

    gr.Examples(
        [
            [
                "It is important for all towns and cities to have large public spaces such as squares and parks. Do you agree or disagree with this statement?",
                (
                    "It is crucial for all metropolitan cities and towns to "
                    "have some recreational facilities like parks and squares because of their numerous benefits. A number of "
                    "arguments surround my opinion, and I will discuss it in upcoming paragraphs. To commence with, the first "
                    "and the foremost merit is that it is beneficial for the health of people because in morning time they can "
                    "go for walking as well as in the evenings, also older people can spend their free time with their loved ones, "
                    "and they can discuss about their daily happenings. In addition, young people do lot of exercise in parks and "
                    "gardens to keep their health fit and healthy, otherwise if there is no park they glue with electronic gadgets "
                    "like mobile phones and computers and many more. Furthermore, little children get best place to play, they play "
                    "with their friends in parks if any garden or square is not available for kids then they use roads and streets "
                    "for playing it can lead to serious incidents. Moreover, parks have some educational value too, in schools, "
                    "students learn about environment protection in their studies and teachers can take their pupils to parks because "
                    "students can see those pictures so lively which they see in their school books and they know about importance "
                    "and protection of trees and flowers. In recapitulate, parks holds immense importance regarding education, health "
                    "for people of every society, so government should build parks in every city and town."
                ),
            ],
        ],
        inputs=[question, essay],
    )

app.launch()