import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient(
    model = "tanmaymahato/llama3-grader-merged-v1",
    token = "your-hf-token",   # add this — needed for your own model
)

def grade_answer(question, reference_answer, student_answer, rubric=""):

    if rubric:
        user_msg = f"""Question: {question}
Rubric: {rubric}
Reference Answer: {reference_answer}
Student Answer: {student_answer}
Please grade the student answer."""
    else:
        user_msg = f"""Question: {question}
Reference Answer: {reference_answer}
Student Answer: {student_answer}
Please grade the student answer."""

    # use chat_completion instead of text_generation
    response = client.chat_completion(
        messages = [
            {
                "role"   : "system",
                "content": "You are an expert grading assistant. Score out of 10 with point-by-point justification."
            },
            {
                "role"   : "user",
                "content": user_msg
            }
        ],
        max_tokens  = 300,
        temperature = 0.1,
    )

    return response.choices[0].message.content

demo = gr.Interface(
    fn      = grade_answer,
    inputs  = [
        gr.Textbox(label="Question",          lines=2),
        gr.Textbox(label="Reference Answer",  lines=3),
        gr.Textbox(label="Student Answer",    lines=3),
        gr.Textbox(label="Rubric (optional)", lines=3),
    ],
    outputs = gr.Textbox(label="Grading Output", lines=10),
    title       = "Llama 3.1 Grading Assistant",
    description = "Fine-tuned Llama 3.1 8B for rubric-based short answer grading — tanmaymahato/llama3-grader-merged-v1",
)

demo.launch()