app.py

import torch
import gradio as gr

from transformers import (
    AutoTokenizer,
    AutoModelForMultipleChoice
)

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

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

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

model.to(device)
model.eval()


# -------------------------
# Prediction Function
# -------------------------
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

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

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

    pred = torch.argmax(probs).item()

    answer = letters[pred]

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

    result += "Scores\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(server_name="0.0.0.0", server_port=7860)