import gradio as gr
from transformers import AutoTokenizer
from inference import load_model, predict

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model, device = load_model("evo_hellaswag.pt")

# Interface logic
def evo_decision(prompt, option1, option2):
    result = predict(model, tokenizer, prompt, option1, option2, device)
    choice = option1 if result["choice"] == 0 else option2
    score_0 = round(result["scores"][0] * 100, 2)
    score_1 = round(result["scores"][1] * 100, 2)
    return (
        f"✅ Evo Suggests: **{choice}**\n\n"
        f"🧠 Confidence Scores:\n"
        f"- Option 1: {score_0}%\n"
        f"- Option 2: {score_1}%"
    )

# UI
with gr.Blocks() as demo:
    gr.Markdown("# 🧬 EvoTransformer – Reasoning API\nAsk Evo a question with 2 choices.")
    
    with gr.Row():
        prompt = gr.Textbox(label="🧠 Scenario or Question", placeholder="e.g. You spilled juice on the floor.")
    with gr.Row():
        option1 = gr.Textbox(label="Option 1", placeholder="Wipe it with a cloth.")
        option2 = gr.Textbox(label="Option 2", placeholder="Ignore and walk away.")
    
    with gr.Row():
        output = gr.Markdown()
    
    with gr.Row():
        btn = gr.Button("Ask Evo")
        btn.click(fn=evo_decision, inputs=[prompt, option1, option2], outputs=[output])

# Launch app
demo.launch()