import spaces
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForMultipleChoice

MODEL_NAME = "Vjay15/electra-mcq"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForMultipleChoice.from_pretrained(MODEL_NAME)
model.eval()

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

EXAMPLES = [
    ["""Prompt: What is the function of mammary glands in mammals?
A: Mammary glands produce milk to feed the young.
B: Mammary glands help mammals draw air into the lungs.
C: Mammary glands help mammals breathe with lungs.
D: Mammary glands excrete nitrogenous waste as urea.
E: Mammary glands separate oxygenated and deoxygenated blood in the mammalian heart."""],
    ["""Prompt: Who was the first to determine the velocity of a star moving away from the Earth using the Doppler effect?
A: Fraunhofer
B: William Huggins
C: Hippolyte Fizeau
D: Vogel and Scheiner
E: None of the above"""],
    ["""Prompt: What is the effect generated by a spinning superconductor?
A: An electric field, precisely aligned with the spin axis.
B: A magnetic field, randomly aligned with the spin axis.
C: A magnetic field, precisely aligned with the spin axis.
D: A gravitational field, randomly aligned with the spin axis.
E: A gravitational field, precisely aligned with the spin axis."""],
    ["""Prompt: What is bollard pull primarily used for measuring?
A: The weight of heavy machinery
B: The speed of locomotives
C: The distance traveled by a truck
D: The strength of tugboats
E: The height of a ballast tractor"""],
]


@spaces.GPU
def predict(text):
    lines = [line.strip() for line in text.strip().splitlines() if line.strip()]

    if len(lines) < 6:
        return "Error", {"message": "Need a prompt line plus 5 option lines (A-E)"}

    question = lines[0].replace("Prompt:", "").strip()

    options = []
    for line in lines[1:6]:
        if ":" not in line:
            return "Error", {"message": f"Option line missing a colon: {line[:40]}"}
        options.append(line.split(":", 1)[1].strip())

    model.to("cuda")

    encoding = tokenizer(
        [question] * 5,
        options,
        truncation=True,
        padding="max_length",
        max_length=256,
        return_tensors="pt",
    )
    inputs = {k: v.unsqueeze(0).to("cuda") for k, v in encoding.items()}

    with torch.no_grad():
        probs = torch.softmax(model(**inputs).logits, dim=1)[0]

    prediction = LABELS[torch.argmax(probs).item()]
    confidence = {label: prob.item() for label, prob in zip(LABELS, probs)}
    return prediction, confidence


demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(
        label="Input",
        lines=10,
        placeholder="""Prompt: What is the capital of France?
A: London
B: Berlin
C: Paris
D: Madrid
E: Rome"""
    ),
    outputs=[
        gr.Textbox(label="Predicted Answer"),
        gr.Label(label="Confidence Scores", num_top_classes=5),
    ],
    examples=EXAMPLES,
    example_labels=[
        "What is the function of mammary glands in mammals?",
        "Who was the first to determine the velocity of a star moving away from the Earth using the Doppler effect?",
        "What is the effect generated by a spinning superconductor?",
        "What is bollard pull primarily used for measuring?",
    ],
    cache_examples=False,
    title="MCQ Solver",
    description="""
A multiple-choice question answering system fine-tuned from Google's ELECTRA Base Discriminator model (`google/electra-base-discriminator`).

Pick an example below, or type your own in this format:

```
Prompt: <question>
A: <option A>
B: <option B>
C: <option C>
D: <option D>
E: <option E>
```
""",
    flagging_mode="never",
)

THEME = gr.themes.Ocean(
    font=["system-ui", "-apple-system", "Segoe UI", "Roboto", "sans-serif"],
    font_mono=["ui-monospace", "SFMono-Regular", "Menlo", "monospace"],
)

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