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

# yaha apna wahi repo_id daalo jo push_to_hub karte waqt use kiya tha
MODEL_NAME = "23f2005181/electra-mcq-solver"  # <- same repo_id jo upar diya tha   # <-- apna username/model-name

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

LETTERS = ['A', 'B', 'C', 'D', 'E']


@spaces.GPU
def predict(prompt, opt_a, opt_b, opt_c, opt_d, opt_e):
    options = [opt_a, opt_b, opt_c, opt_d, opt_e]

    if not prompt.strip() or any(not o.strip() for o in options):
        return "Please prompt aur saare 5 options fill karo."

    prompts = [prompt] * 5

    # tokenizer 5 (prompt, option) pairs banata hai -- yehi
    # AutoModelForMultipleChoice ka expected input format hai
    encoding = tokenizer(
        prompts, options,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=160,
    )

    # batch dimension add karo: (5, seq_len) -> (1, 5, seq_len)
    inputs = {k: v.unsqueeze(0) for k, v in encoding.items()}

    with torch.no_grad():
        outputs = model(**inputs)

    logits = outputs.logits[0]                     # shape: (5,)
    probs = torch.softmax(logits, dim=0).tolist()

    ranked = sorted(zip(LETTERS, probs), key=lambda x: -x[1])
    top3 = ranked[:3]

    lines = [f"{i+1}. Option {letter}  —  {prob*100:.1f}% confidence"
              for i, (letter, prob) in enumerate(top3)]
    return "\n".join(lines)


demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(label="Question / Prompt", lines=2, placeholder="e.g. Which of the following best describes..."),
        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="Top-3 Predicted Answers", lines=4),
    title="Smart MCQ Solver — ELECTRA (fine-tuned)",
    description=(
        "Multiple-choice question daalo (prompt + 5 options), model top-3 "
        "sabse likely correct answers ranked confidence ke saath dega. "
        "Fine-tuned ELECTRA-base-discriminator model, Hugging Face "
        "AutoModelForMultipleChoice head ke saath."
    ),
    examples=[
        [
            "Which force is responsible for keeping planets in orbit around the sun?",
            "Electromagnetic force",
            "Gravitational force",
            "Nuclear force",
            "Frictional force",
            "Centripetal force alone",
        ]
    ],
)

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