import torch
import gradio as gr
import spaces

from transformers import (
    AutoTokenizer,
    AutoModelForMultipleChoice
)


# -------------------------
# Model
# -------------------------

MODEL_PATH = "./best_model_fold_1"


# -------------------------
# Device
# -------------------------

device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)

print("Device:", device)


# -------------------------
# Load tokenizer
# -------------------------

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_PATH
)

print("Tokenizer loaded successfully!")


# -------------------------
# Load model
# -------------------------

model = AutoModelForMultipleChoice.from_pretrained(
    MODEL_PATH
)

model.to(device)
model.eval()

print("Model loaded successfully!")


# -------------------------
# Prediction
# -------------------------
@spaces.GPU(duration=30)
def predict(question, A, B, C, D, E):

    choices = [A, B, C, D, E]

    encoding = tokenizer(
        [question] * 5,
        choices,
        max_length=256,
        truncation=True,
        padding="max_length",
        return_tensors="pt"
    )

    input_ids = encoding["input_ids"].unsqueeze(0).to(device)

    attention_mask = encoding[
        "attention_mask"
    ].unsqueeze(0).to(device)

    with torch.no_grad():

        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask
        )

    logits = outputs.logits

    probabilities = torch.softmax(
        logits,
        dim=1
    )[0]

    letters = ["A", "B", "C", "D", "E"]

    prediction = torch.argmax(
        probabilities
    ).item()

    answer = letters[prediction]

    result = (
        f"Predicted Answer: {answer}\n\n"
        "Scores\n\n"
    )

    for letter, score in zip(
        letters,
        probabilities
    ):
        result += (
            f"{letter}: "
            f"{score.item():.4f}\n"
        )

    return result


# -------------------------
# Gradio
# -------------------------

demo = gr.Interface(

    fn=predict,

    inputs=[

        gr.Textbox(
            lines=4,
            label="Question"
        ),

        gr.Textbox(
            label="Option A"
        ),

        gr.Textbox(
            label="Option B"
        ),

        gr.Textbox(
            label="Option C"
        ),

        gr.Textbox(
            label="Option D"
        ),

        gr.Textbox(
            label="Option E"
        ),
    ],

    outputs=gr.Textbox(
        label="Prediction"
    ),

    title="DeBERTa MCQ Solver",

    description=(
        "Enter a question and five options."
    ),

    examples=[
        [
            "Which planet is known as the Red Planet?",
            "Earth",
            "Mars",
            "Venus",
            "Jupiter",
            "Saturn"
        ]
    ]
)


if __name__ == "__main__":
    demo.launch()