import torch
import gradio as gr
import spaces

from transformers import (
    AutoTokenizer,
    AutoModelForMultipleChoice
)



# -------------------------
# Load tokenizer
# -------------------------
MODEL_NAME = "rohitk123/roberta-base-mcq-solver"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
# -------------------------
# Load model
# -------------------------
model = AutoModelForMultipleChoice.from_pretrained(MODEL_NAME)
model.eval()



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

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

    model.to(device)

    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
    probs = torch.softmax(logits, dim=1)[0]

    letters = ["A", "B", "C", "D", "E"]
    pred = torch.argmax(probs).item()

    result = f"Predicted Answer: {letters[pred]}\n\nScores\n\n"

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

    return result


# -------------------------
# Gradio Interface
# -------------------------
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="RoBERTa-Base 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()