app.py

import gradio as gr
import torch
from functools import lru_cache
from transformers import AutoTokenizer, AutoModelForCausalLM

try:
    import spaces
except Exception:
    class _SpacesFallback:
        @staticmethod
        def GPU(fn):
            return fn
    spaces = _SpacesFallback()


MODEL_CHOICES = {
    "SmolLM2 135M Instruct": "HuggingFaceTB/SmolLM2-135M-Instruct",
    "SmolLM2 360M Instruct": "HuggingFaceTB/SmolLM2-360M-Instruct",
    "SmolLM2 1.7B Instruct": "HuggingFaceTB/SmolLM2-1.7B-Instruct",
}

DEFAULT_MODEL_LABEL = "SmolLM2 360M Instruct"
DEFAULT_BACKEND = "GPU"

CSS = """
#app-title {
    font-size: 2.2rem !important;
    font-weight: 900 !important;
    margin-bottom: 0.15rem !important;
    letter-spacing: -0.03em;
}
#app-subtitle {
    color: #6b7280;
    margin-bottom: 1rem;
    font-size: 1rem;
}
.control-card {
    border: 1px solid rgba(128,128,128,0.18);
    border-radius: 20px;
    padding: 16px;
    background: linear-gradient(180deg, rgba(255,255,255,0.06), rgba(255,255,255,0.03));
    box-shadow: 0 8px 24px rgba(0,0,0,0.08);
}
.gradio-container {
    max-width: 1200px !important;
}
"""


def _get_device(requested: str) -> str:
    if requested == "cuda" and torch.cuda.is_available():
        return "cuda"
    return "cpu"


def _model_device(model) -> torch.device:
    for param in model.parameters():
        if param.device.type != "meta":
            return param.device
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


@lru_cache(maxsize=8)
def load_model(model_id: str, device: str):
    device = _get_device(device)

    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    if device == "cuda":
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            dtype=torch.float16,
            device_map="auto",
            low_cpu_mem_usage=True,
        )
    else:
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            dtype=torch.float32,
            low_cpu_mem_usage=True,
        ).to(device)

    model.eval()
    return tokenizer, model


def _extract_text(content) -> str:
    """Gradio 6's messages format stores `content` as either a plain string
    or a list of content blocks (e.g. [{"type": "text", "text": "..."}]).
    Normalize either shape down to plain text."""
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for block in content:
            if isinstance(block, str):
                parts.append(block)
            elif isinstance(block, dict):
                if "text" in block:
                    parts.append(block.get("text") or "")
        return "".join(parts)
    if isinstance(content, dict):
        return content.get("text", "") or ""
    return str(content)


def build_messages(history, user_message: str):
    """history is a list of {'role': ..., 'content': ...} dicts
    (Gradio Chatbot 'messages' format). content may be a string or a
    list of content blocks."""
    messages = []
    for msg in history or []:
        role = msg.get("role")
        text = _extract_text(msg.get("content"))
        if role in ("user", "assistant") and text:
            messages.append({"role": role, "content": text})

    messages.append({"role": "user", "content": user_message})
    return messages


def build_fallback_prompt(history, user_message: str) -> str:
    parts = []
    for msg in history or []:
        role = msg.get("role")
        text = _extract_text(msg.get("content"))
        if not text:
            continue
        if role == "user":
            parts.append(f"User: {text}")
        elif role == "assistant":
            parts.append(f"Assistant: {text}")

    parts.append(f"User: {user_message}")
    parts.append("Assistant:")
    return "\n".join(parts)


def _generate(message, history, model_label, max_new_tokens, device_name):
    model_id = MODEL_CHOICES[model_label]
    tokenizer, model = load_model(model_id, device_name)

    max_new_tokens = int(max_new_tokens)

    if getattr(tokenizer, "chat_template", None):
        messages = build_messages(history, message)
        encoded = tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            return_tensors="pt",
            return_dict=True,
        )
        input_ids = encoded["input_ids"]
    else:
        prompt = build_fallback_prompt(history, message)
        input_ids = tokenizer(prompt, return_tensors="pt").input_ids

    input_ids = input_ids.to(_model_device(model))
    attention_mask = torch.ones_like(input_ids)

    generation_kwargs = {
        "attention_mask": attention_mask,
        "max_new_tokens": max_new_tokens,
        "do_sample": True,
        "temperature": 0.7,
        "top_p": 0.95,
        "eos_token_id": tokenizer.eos_token_id,
        "pad_token_id": tokenizer.eos_token_id,
    }

    with torch.inference_mode():
        output_ids = model.generate(input_ids, **generation_kwargs)

    new_tokens = output_ids[0][input_ids.shape[-1]:]
    reply = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
    return reply or ""


def generate_cpu(message, history, model_label, max_new_tokens):
    return _generate(message, history, model_label, max_new_tokens, device_name="cpu")


@spaces.GPU
def generate_gpu(message, history, model_label, max_new_tokens):
    return _generate(message, history, model_label, max_new_tokens, device_name="cuda")


def respond(message, history, model_label, backend, max_new_tokens):
    message = (message or "").strip()
    history = history or []

    if not message:
        return history, ""

    if backend == "GPU":
        reply = generate_gpu(message, history, model_label, max_new_tokens)
    else:
        reply = generate_cpu(message, history, model_label, max_new_tokens)

    history = history + [
        {"role": "user", "content": message},
        {"role": "assistant", "content": reply},
    ]
    return history, ""


with gr.Blocks(title="SmolLM2 Chat") as demo:
    gr.Markdown("# SmolLM2 Chat", elem_id="app-title")
    gr.Markdown(
        "A simple chat interface for SmolLM2 instruct models.",
        elem_id="app-subtitle",
    )

    with gr.Row():
        with gr.Column(scale=1, min_width=300):
            gr.Markdown("### Controls")
            with gr.Group(elem_classes=["control-card"]):
                model_dropdown = gr.Dropdown(
                    choices=list(MODEL_CHOICES.keys()),
                    value=DEFAULT_MODEL_LABEL,
                    label="Model",
                    interactive=True,
                )
                backend = gr.Radio(
                    choices=["CPU", "GPU"],
                    value=DEFAULT_BACKEND,
                    label="Backend",
                    interactive=True,
                )
                max_new_tokens = gr.Slider(
                    minimum=32,
                    maximum=2048,
                    value=256,
                    step=1,
                    label="Max new tokens",
                )

        with gr.Column(scale=2, min_width=500):
            gr.Markdown("### Chat")
            chatbot = gr.Chatbot(
                height=620,
                label="Conversation",
            )
            message = gr.Textbox(
                placeholder="PauliePocket is just better than you...",
                label="Message",
                lines=3,
            )

            gr.Examples(
                examples=[
                    ["Explain transformers in ML."],
                    ["Write a story about the 2026 World Cup."],
                    ["Generate a song about a girl falling in love."],
                ],
                inputs=message,
                label="Try these stupid prompts",
            )

            with gr.Row():
                send = gr.Button("Send", variant="primary")
                clear = gr.Button("Clear")

    send.click(
        respond,
        inputs=[message, chatbot, model_dropdown, backend, max_new_tokens],
        outputs=[chatbot, message],
    )
    message.submit(
        respond,
        inputs=[message, chatbot, model_dropdown, backend, max_new_tokens],
        outputs=[chatbot, message],
    )
    clear.click(lambda: ([], ""), outputs=[chatbot, message])

if __name__ == "__main__":
    demo.launch(css=CSS, theme=gr.themes.Soft())