app.py

"""Public Gradio showcase for NPCMicro13M.

The visible chat transcript is for the visitor.  The model itself receives only
the current villager state and the current player message on every turn.
"""

from __future__ import annotations

import sys
import time
from functools import lru_cache
import inspect
from pathlib import Path

# Hugging Face provides this package in ZeroGPU Spaces.  The fallback keeps
# local development possible on machines without the Spaces runtime.
try:
    import spaces
except ImportError:  # pragma: no cover - local development only
    class _LocalSpaces:
        @staticmethod
        def GPU(function=None, **_kwargs):
            if function is None:
                return lambda wrapped: wrapped
            return function

    spaces = _LocalSpaces()

import gradio as gr
from tokenizers import Tokenizer


ROOT = Path(__file__).resolve().parent
DEPLOYMENT = ROOT / "deployment"
sys.path.insert(0, str(DEPLOYMENT))

from uomind_grounded_response import grounded_response  # noqa: E402
from uomind_infer import (  # noqa: E402
    choose_device,
    choose_precision,
    discover_checkpoint,
    discover_tokenizer,
    generate,
    load_model,
    validate_tokenizer,
)

_GRADIO_MAJOR = int(str(getattr(gr, "__version__", "0")).split(".", 1)[0] or 0)
# Gradio 6 removed tuple histories.  Prefer the explicit messages mode even
# when a wrapper hides the Chatbot constructor signature.
_CHATBOT_ACCEPTS_TYPE = "type" in inspect.signature(gr.Chatbot).parameters
CHAT_SUPPORTS_MESSAGES = (
    _GRADIO_MAJOR >= 6 or _CHATBOT_ACCEPTS_TYPE
)


def article_for(role: str) -> str:
    return "an" if role[:1].lower() in "aeiou" else "a"


def clean_field(value: str, fallback: str) -> str:
    value = " ".join(str(value or "").replace("\n", " ").split()).strip()
    return value or fallback


def build_state(name: str, profession: str, home: str, details: str) -> str:
    name = clean_field(name, "Sella")
    profession = clean_field(profession, "baker")
    home = clean_field(home, "Britain")
    details = " ".join(str(details or "").replace("\n", " ").split()).strip()
    state = f"Your name is {name}. You are {article_for(profession)} {profession} from {home}."
    if details:
        state += f" {details.rstrip('.')}."
    return state


@lru_cache(maxsize=1)
def runtime():
    """Load the compact model once per Space process."""
    tokenizer_path = discover_tokenizer(ROOT, None)
    tokenizer = Tokenizer.from_file(str(tokenizer_path))
    validate_tokenizer(tokenizer)
    # ZeroGPU gives this process a CUDA device only while answer() is running.
    device = choose_device("auto")
    precision = choose_precision(device, "auto")
    checkpoint = discover_checkpoint(ROOT, None)
    model, _payload, _config = load_model(
        checkpoint,
        tokenizer,
        device,
        precision,
        attention_backend=None,
    )
    return model, tokenizer, device, precision


@spaces.GPU(duration=60)
def answer(
    message: str,
    history: list[dict],
    name: str,
    profession: str,
    home: str,
    details: str,
    protect_facts: bool,
    temperature: float,
):
    message = " ".join(str(message or "").split()).strip()
    history = list(history or [])
    if not message:
        return history, "", "Write a message for the villager first."

    state = build_state(name, profession, home, details)
    model, tokenizer, device, precision = runtime()
    started = time.perf_counter()
    result = generate(
        model,
        tokenizer,
        state=state,
        player=message,
        device=device,
        precision=precision,
        max_new_tokens=32,
        temperature=float(temperature),
        top_p=0.9,
        max_system_tokens=40,
    )
    raw = result["text"] or "I know not what to say."
    final = (
        grounded_response(state, message, lambda _state, _player: raw)
        if protect_facts
        else raw
    )
    elapsed = time.perf_counter() - started
    if CHAT_SUPPORTS_MESSAGES:
        history.extend(
            [
                {"role": "user", "content": message},
                {"role": "assistant", "content": final},
            ]
        )
    else:
        history.append([message, final])
    layer_note = "identity/profession protection on" if protect_facts else "raw model output"
    if final != raw and protect_facts:
        layer_note += " • protection correction applied"
    status = (
        f"{layer_note} • {str(device).upper()}/{str(precision).replace('torch.', '').upper()} • "
        f"{elapsed * 1000:.0f} ms • each turn is stateless"
    )
    return history, "", status


def reset_chat():
    return [], "", "Ready for a new conversation."


def reset_villager():
    return [], "", "Sella", "baker", "Britain", "Friendly, patient, and familiar with the local market."


DESCRIPTION = """
## Meet a villager

NPCMicro13M is a compact, CPU-friendly language model trained for everyday
villager roles: bakers, blacksmiths, farmers, merchants, innkeepers, healers,
craftsmen, and more. Define the villager once, then ask about their work,
home, local knowledge, wares, weather, or daily life.

The transcript is displayed for you, but the model receives only the villager
state and your current message on each turn. This keeps the NPC from trying to
continue an old conversation. Quests, combat, movement, and game rules belong
in the game server, not in this dialogue model.
"""


with gr.Blocks(
    title="NPCMicro13M Villager Chat",
) as demo:
    with gr.Column(elem_classes="hero"):
        gr.Markdown(DESCRIPTION)

        with gr.Row():
            with gr.Column(scale=1, elem_classes="state-panel"):
                gr.Markdown("### Villager details")
                name = gr.Textbox(value="Sella", label="Name")
                profession = gr.Textbox(value="baker", label="Profession")
                home = gr.Textbox(value="Britain", label="Home town")
                details = gr.Textbox(
                    value="Friendly, patient, and familiar with the local market.",
                    label="Personality and local knowledge",
                    lines=3,
                )
                new_villager = gr.Button("Reset to example villager", variant="secondary")
                protect_facts = gr.Checkbox(
                    value=True,
                    label="Protect identity and profession consistency",
                    info="Contradictory self-claims are corrected; local topics go to the model.",
                )
                temperature = gr.Slider(
                    minimum=0.0,
                    maximum=0.8,
                    value=0.25,
                    step=0.05,
                    label="Creativity",
                    info="Lower is more consistent; higher is more varied.",
                )

            with gr.Column(scale=2):
                chat_options = dict(
                    label="Conversation",
                    height=520,
                    placeholder="Ask the villager about their work or town...",
                )
                if _CHATBOT_ACCEPTS_TYPE:
                    chat_options["type"] = "messages"
                chat = gr.Chatbot(**chat_options)
                with gr.Row():
                    message = gr.Textbox(
                        label="You",
                        placeholder="What do you sell today?",
                        scale=5,
                        show_label=False,
                    )
                    send = gr.Button("Send", variant="primary", scale=1)
                clear = gr.Button("Clear conversation", variant="secondary")
                status = gr.Markdown("Ready. The model will load on the first message.", elem_classes="status")

        gr.Examples(
            examples=[
                ["What do you sell today?"],
                ["How long have you lived in this town?"],
                ["Can you tell me about your work?"],
                ["Is the road to the market busy this morning?"],
                ["Have a good day."],
            ],
            inputs=message,
            label="Try a question",
        )

    inputs = [message, chat, name, profession, home, details, protect_facts, temperature]
    outputs = [chat, message, status]
    send.click(answer, inputs=inputs, outputs=outputs)
    message.submit(answer, inputs=inputs, outputs=outputs)
    clear.click(reset_chat, outputs=outputs)
    new_villager.click(
        reset_villager,
        outputs=[chat, message, name, profession, home, details],
    )


if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Soft(primary_hue="amber", neutral_hue="stone"),
        css="""
        .hero { max-width: 980px; margin: 0 auto; }
        .state-panel { border: 1px solid #d7c6a5; border-radius: 12px; padding: 10px; }
        .status { color: #6b6255; font-size: 0.88rem; }
        footer { display: none !important; }
        """,
    )