"""
Queensland AI Meetup — Base vs Fine-Tuned Gemma 3 270M

Side-by-side comparison: same input goes to both models.
Shows what 2 epochs of fine-tuning on ~7,700 training pairs does to a 270M parameter model.

Deployed to: mrdbourke/queensland-ai-fine-tuned-gemma3
"""

import time
import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# ── Load both models ────────────────────────────────────────────────────

BASE_MODEL_NAME = "google/gemma-3-270m-it"
FINE_TUNED_MODEL_NAME = "mrdbourke/queensland-ai-gemma3-fine-tuned"  # Update if your repo name differs

print("[INFO] Loading base model...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_NAME,
    dtype="auto",
    device_map="auto",
    attn_implementation="eager",
)
base_tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME)
base_pipe = pipeline("text-generation", model=base_model, tokenizer=base_tokenizer)

print("[INFO] Loading fine-tuned model...")
ft_model = AutoModelForCausalLM.from_pretrained(
    FINE_TUNED_MODEL_NAME,
    dtype="auto",
    device_map="auto",
    attn_implementation="eager",
)
ft_tokenizer = AutoTokenizer.from_pretrained(FINE_TUNED_MODEL_NAME)
ft_pipe = pipeline("text-generation", model=ft_model, tokenizer=ft_tokenizer)

print("[INFO] Both models loaded!")


# ── Inference ───────────────────────────────────────────────────────────

def run_model(pipe, input_text, max_new_tokens=512):
    start_time = time.time()
    raw_output = pipe(
        text_inputs=[{"role": "user", "content": input_text}],
        max_new_tokens=max_new_tokens,
        disable_compile=True,
    )
    elapsed = round(time.time() - start_time, 4)
    generated = raw_output[0]["generated_text"]
    output_text = generated[-1]["content"] if isinstance(generated, list) else str(generated)
    return output_text, elapsed


@spaces.GPU
def compare(input_text):
    base_output, base_time = run_model(base_pipe, input_text)
    ft_output, ft_time = run_model(ft_pipe, input_text)
    return base_output, base_time, ft_output, ft_time


# ── Gradio UI ───────────────────────────────────────────────────────────

description = """## Base vs Fine-Tuned Gemma 3 270M — Side by Side

Type a name or question and see both models respond. The **base model** (left) has no knowledge of meetup attendees. The **fine-tuned model** (right) was trained on ~7,700 examples for 2 epochs.

**Try these:**
- An attendee: `Daniel Bourke`, `Adam Luff`, `Crighton Nichols`
- Lowercase: `daniel bourke`, `william bourke`
- A non-attendee: `Elon Musk`, `Shrek`, `Pikachu`
- A short name: `Leah H.`, `chrism`
"""

with gr.Blocks(title="Queensland AI Meetup — Base vs Fine-Tuned") as demo:
    gr.Markdown("# 🧠 Queensland AI Meetup — Base vs Fine-Tuned Gemma 3 270M")
    gr.Markdown(description)

    with gr.Row():
        input_text = gr.Textbox(
            lines=2,
            label="Type a name or question",
            placeholder="e.g. Daniel Bourke",
            scale=3,
        )
        submit_btn = gr.Button("Compare", variant="primary", scale=1)

    with gr.Row():
        with gr.Column():
            gr.Markdown("### ❌ Base Model (google/gemma-3-270m-it)")
            base_output = gr.Textbox(lines=12, label="Base Model Output", interactive=False)
            base_time = gr.Number(label="Base Time (s)")

        with gr.Column():
            gr.Markdown("### ✅ Fine-Tuned Model")
            ft_output = gr.Textbox(lines=12, label="Fine-Tuned Output", interactive=False)
            ft_time = gr.Number(label="Fine-Tuned Time (s)")

    submit_btn.click(
        fn=compare,
        inputs=[input_text],
        outputs=[base_output, base_time, ft_output, ft_time],
    )

    input_text.submit(
        fn=compare,
        inputs=[input_text],
        outputs=[base_output, base_time, ft_output, ft_time],
    )

    gr.Examples(
        examples=[
            ["Daniel Bourke"],
            ["daniel bourke"],
            ["William Bourke"],
            ["What are Daniel Bourke's skills?"],
            ["Adam Luff"],
            ["Elon Musk"],
            ["Shrek"],
            ["Leah H."],
            ["Crighton Nichols"],
        ],
        inputs=input_text,
    )

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