app.py

"""BytesTalk — PersonaMini Chat.  A public Gradio Space for the three PersonaMini-1 models.



Everything is per-session. Gradio's gr.State lives in the browser session, so two people using the

Space at the same time never see each other's conversation, and a refresh starts clean - nothing is

written to disk and nothing is shared.



Models are pulled from the Hub at first use and cached in the Space:



    bytestalkai/PersonaMini-1-small     28.8M   GPT-2 style export, 256-token context

    bytestalkai/PersonaMini-1-medium    63.2M   custom code, 512-token context

    bytestalkai/PersonaMini-1-big      160.0M   custom code, 1024-token context



ZeroGPU: generation runs inside a @spaces.GPU function, so the GPU is only held for the duration

of a reply.

"""
import os
import re
import json
import threading
import hmac

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

try:
    import spaces                       # only present on a ZeroGPU Space
    HAS_ZERO = True
except ImportError:                     # running locally
    HAS_ZERO = False

    class spaces:                       # noqa: N801  - shim so the decorator still works
        @staticmethod
        def GPU(*a, **kw):
            def deco(fn):
                return fn
            return deco if not a or not callable(a[0]) else a[0]


HERE = os.path.dirname(os.path.abspath(__file__))
EOT = 50256

MODELS = {
    "PersonaMini-1-big · 160M": dict(
        repo="bytestalkai/PersonaMini-1-big", ctx=1024, custom=True,
        note="Newest and strongest — explicit roleplay, holds character, refuses the hard lines."),
    "PersonaMini-1-medium · 63.2M": dict(
        repo="bytestalkai/PersonaMini-1-medium", ctx=512, custom=True,
        note="Good memory and identity for its size."),
    "PersonaMini-1-small · 28.8M": dict(
        repo="bytestalkai/PersonaMini-1-small", ctx=256, custom=False,
        note="The first release. Fast, but loses the thread quickly."),
}
DEFAULT = "PersonaMini-1-big · 160M"

PRESETS = {
    # Keep sampling open, as in the original PersonaMini UI.  The model was
    # trained with repetition allowed; forcing penalties makes its phrasing
    # noticeably less natural, especially in songs and roleplay.
    "Default":  dict(temp=0.85, top_k=50, top_p=1.00, rep=1.00, no_repeat=0),
    "Precise":  dict(temp=0.60, top_k=40, top_p=0.95, rep=1.00, no_repeat=0),
    "Creative": dict(temp=1.00, top_k=60, top_p=1.00, rep=1.00, no_repeat=0),
    "Wild":     dict(temp=1.20, top_k=80, top_p=1.00, rep=1.00, no_repeat=0),
}

_loaded = {}
_load_lock = threading.Lock()


def get_model(name):
    """Load once per process. The Space keeps whichever models people actually use."""
    with _load_lock:
        if name in _loaded:
            return _loaded[name]
        spec = MODELS[name]
        tok = AutoTokenizer.from_pretrained(spec["repo"])
        m = AutoModelForCausalLM.from_pretrained(
            spec["repo"], trust_remote_code=spec["custom"])
        m = m.eval()
        _loaded[name] = (m, tok, spec)
        return _loaded[name]


# --------------------------------------------------------------------------- characters
def load_characters():
    try:
        return json.load(open(os.path.join(HERE, "characters.json"), encoding="utf-8"))["characters"]
    except (OSError, ValueError, KeyError):
        return []


CHARS = load_characters()
BY_NAME = {c["name"]: c for c in CHARS}


def avatar_path(c):
    for cand in (c.get("avatar") or "", c["id"] + ".png", c["name"].split()[0].lower() + ".png"):
        p = os.path.join(HERE, "assets", cand)
        if cand and os.path.exists(p):
            return p
    return None


# --------------------------------------------------------------------------- tiny RAG
_WORD = re.compile(r"[a-z0-9']+")
_STOP = set("""a an and are as at be but by for from had has have he her his i if in is it its me

my no not of on or our she so that the their them then there they this to was we were what when

which who will with you your just like get got do does did dont im ive youre about all can cant""".split())


def _bag(t):
    return {w for w in _WORD.findall(t.lower()) if w not in _STOP and len(w) > 2}


def recall(history, card, keep_recent=6):
    """Pull a couple of relevant older lines back into context.



    The context window is 256-1024 tokens depending on the model, so in any real conversation the

    oldest turns fall out and the model forgets what it was told. Plain word overlap, no embedding

    model - it catches names and stated facts, which is most of what people notice going missing.

    """
    if len(history) <= keep_recent:
        older = []
    else:
        older = history[:-keep_recent]
    last = history[-1][0] if history else ""
    docs = []
    for u, a in older:
        if u and len(u) > 20:
            docs.append(f"You said earlier: {u.strip()[:190]}")
    for ln in (card or "").splitlines():
        if len(ln.strip()) > 25:
            docs.append(ln.strip()[:190])
    q = _bag(last)
    if not q or not docs:
        return []
    df = {}
    for d in docs:
        for w in _bag(d):
            df[w] = df.get(w, 0) + 1
    n = len(docs)
    scored = []
    for d in docs:
        b = _bag(d)
        if b:
            s = sum(1.0 / (1 + df.get(w, 0) / n) for w in (q & b)) / (1 + len(b) ** 0.4)
            if s > 0:
                scored.append((s, d))
    scored.sort(key=lambda x: -x[0])
    return [d for _, d in scored[:3]]


def build_prompt(card, memory, history, ctx, reserve, tok):
    head = ""
    if card.strip():
        head += card.strip() + "\n\n"
    facts = [f for f in (memory or []) if f.strip()]
    if facts:
        head += "[MEMORY]\n" + "\n".join("- " + f for f in facts) + "\n[/MEMORY]\n\n"
    head_ids = tok(head).input_ids if head else []
    tail_ids = tok("### ASSISTANT:\n").input_ids

    turns = []
    for u, a in history:
        if u:
            turns.append(tok(f"### USER:\n{u.strip()}\n\n").input_ids)
        if a:
            turns.append(tok(f"### ASSISTANT:\n{a.strip()}\n\n").input_ids)
    budget = ctx - reserve - len(head_ids) - len(tail_ids)
    while turns and sum(len(t) for t in turns) > budget:
        turns.pop(0)                      # oldest first; the card always survives
    ids = head_ids + [t for turn in turns for t in turn] + tail_ids
    return ids[-(ctx - reserve):]


STOPS = ("### USER:", "\n### ", "### ASSISTANT:")
LONG_WORDS = ("song", "lyric", "poem", "screenplay", "scene", "script", "recipe", "story",
              "html", "page", "code", "essay", "letter", "list")


def text_content(value):
    """Convert Gradio's text/multimodal message shapes into plain prompt text."""
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        return text_content(value.get("text", value.get("content", "")))
    if isinstance(value, list):
        return "\n".join(part for part in (text_content(item) for item in value) if part)
    return "" if value is None else str(value)


@spaces.GPU(duration=60)
@torch.no_grad()
def generate_stream(name, ids, temp, top_k, top_p, rep, no_repeat, max_new, min_new):
    """Yield the reply as it is sampled so both Gradio and the website stream it."""
    m, tok, spec = get_model(name)
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    m = m.to(dev)
    ctx = torch.tensor([ids], device=dev)
    seq, out = list(ids), []
    previous = ""
    for _ in range(max_new):
        res = m(input_ids=ctx[:, -spec["ctx"]:])
        logits = (res.logits if hasattr(res, "logits") else res[0])[0, -1].float().clone()
        if rep and rep != 1.0:
            idx = torch.tensor(list(set(seq[-256:])), device=dev)
            v = logits[idx]
            logits[idx] = torch.where(v > 0, v / rep, v * rep)
        if no_repeat and len(seq) >= no_repeat - 1:
            prefix = tuple(seq[-(no_repeat - 1):])
            for i in range(len(seq) - no_repeat + 1):
                if tuple(seq[i:i + no_repeat - 1]) == prefix:
                    logits[seq[i + no_repeat - 1]] = -1e10
        logits = logits / max(temp, 1e-5)
        kth = torch.topk(logits, min(top_k, logits.numel()))[0][-1]
        logits = logits.masked_fill(logits < kth, float("-inf"))
        probs = torch.softmax(logits, -1)
        if top_p and top_p < 1.0:
            sorted_probs, sorted_ids = torch.sort(probs, descending=True)
            cut = torch.cumsum(sorted_probs, 0) - sorted_probs > top_p
            sorted_probs[cut] = 0.0
            sorted_probs = sorted_probs / sorted_probs.sum()
            nxt = sorted_ids[torch.multinomial(sorted_probs, 1)]
        else:
            nxt = torch.multinomial(probs, 1)
        t = int(nxt.item())
        if t == EOT:
            if len(out) < min_new:
                continue                  # too early to stop on a "write me a song"
            break
        out.append(t)
        seq.append(t)
        ctx = torch.cat([ctx, nxt.view(1, 1)], 1)
        txt = tok.decode(out).replace("\ufffd", "")
        # Token decoding can occasionally revise the final character of the
        # previous piece, so clients receive the complete text-so-far.
        if txt != previous:
            previous = txt
            yield txt
        if any(s in txt for s in STOPS):
            break
    txt = tok.decode(out).replace("\ufffd", "")
    for s in STOPS:
        txt = txt.split(s)[0]
    if txt.strip() != previous.strip():
        yield txt.strip()


def generate(name, ids, temp, top_k, top_p, rep, no_repeat, max_new, min_new):
    """Compatibility wrapper for the normal browser chat UI."""
    final = ""
    for final in generate_stream(name, ids, temp, top_k, top_p, rep, no_repeat, max_new, min_new):
        pass
    return final.strip()


# --------------------------------------------------------------------------- chat
def respond(message, history, model_name, preset, card, memory_txt):
    """history is Gradio's messages-format list; state lives in the browser session only."""
    message = text_content(message).strip()
    if not message:
        return history, ""

    pairs = []
    u = None
    for h in history:
        if h["role"] == "user":
            u = text_content(h.get("content"))
        elif u is not None:
            pairs.append((u, text_content(h.get("content"))))
            u = None
    pairs.append((message, ""))

    facts = [ln.strip() for ln in (memory_txt or "").splitlines() if ln.strip()]
    facts += recall(pairs, card)

    _m, tok, spec = get_model(model_name)
    max_new = max(200, spec["ctx"] // 2)
    reserve = min(max_new + 8, int(spec["ctx"] * 0.6))
    ids = build_prompt(card or "", facts, pairs, spec["ctx"], reserve, tok)

    low = message.lower()
    n_lines = re.search(r"\b(\d+)\s*(lines?|words?|sentences?|items?)\b", low)
    min_new = 0 if (n_lines and int(n_lines.group(1)) <= 6) else (
        160 if any(w in low for w in LONG_WORDS) else 0)

    p = PRESETS.get(preset, PRESETS["Default"])
    reply = generate(model_name, ids, p["temp"], p["top_k"], p["top_p"], p["rep"],
                     p["no_repeat"], max_new, min_new)

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


def start_character(name, history):
    """Load a card and open with the character's greeting, so the scene has already begun."""
    c = BY_NAME.get(name)
    if not c:
        return history, "", gr.update()
    greet = c.get("greeting") or ""
    new_hist = [{"role": "assistant", "content": greet}] if greet else []
    return new_hist, c.get("card_text") or "", gr.update(value=f"### {c['name']}\n{c.get('tagline','')}")


CSS = """

.gradio-container{max-width:1180px !important}

#title h1{font-size:26px;margin-bottom:2px}

#title p{color:#a99aa0;margin-top:0}

.charcard{border:1px solid #392a3c;border-radius:12px;padding:8px;background:#211826}

footer{display:none !important}

"""

# gradio 6 moved theme and css from the Blocks constructor to launch()
with gr.Blocks(title="BytesTalk — PersonaMini Chat") as demo:
    gr.Markdown(
        "# BytesTalk — PersonaMini Chat\n"
        "Three roleplay language models trained **from scratch** by BytesTalk — 28.8M, 63.2M and "
        "160M parameters. Nothing is stored: your conversation lives in your browser session only "
        "and disappears when you refresh.\n\n"
        "**18+.** These models write adult fiction. They refuse sexual content involving minors, "
        "non-consent, and bestiality. They are tiny — they get facts wrong constantly.",
        elem_id="title")

    with gr.Row():
        with gr.Column(scale=3):
            chat = gr.Chatbot(height=470, show_label=False)
            with gr.Row():
                msg = gr.Textbox(placeholder="Message PersonaMini…", scale=8,
                                 show_label=False, autofocus=True, lines=1, max_lines=6)
                send = gr.Button("Send", variant="primary", scale=1, min_width=90)
            with gr.Row():
                clear = gr.Button("New chat", size="sm")
                retry = gr.Button("Regenerate", size="sm")
                undo = gr.Button("Undo", size="sm")

        with gr.Column(scale=2):
            model_dd = gr.Dropdown(list(MODELS), value=DEFAULT, label="Model")
            model_note = gr.Markdown(MODELS[DEFAULT]["note"])
            preset_dd = gr.Dropdown(list(PRESETS), value="Default", label="Style")

            with gr.Accordion("Character card", open=False):
                card_box = gr.Textbox(lines=6, show_label=False,
                                      placeholder="Character: Mira (she/her)\n"
                                                  "A deadpan roommate with dry humour.")
                mem_box = gr.Textbox(lines=3, label="Things it should remember about you",
                                     placeholder="My name is Ali.")

            with gr.Accordion("Character library", open=True):
                char_dd = gr.Dropdown([c["name"] for c in CHARS],
                                      label="Pick someone to talk to",
                                      value=None, filterable=True)
                char_info = gr.Markdown("")
                char_img = gr.Image(height=190, show_label=False, container=False)
                start_btn = gr.Button("Start chat with them", variant="secondary")

    def on_model(n):
        return MODELS[n]["note"]

    def on_pick(n):
        c = BY_NAME.get(n)
        if not c:
            return None, ""
        return avatar_path(c), f"**{c['name']}** — *{c.get('media','')}*\n\n{c.get('tagline','')}"

    model_dd.change(on_model, model_dd, model_note)
    char_dd.change(on_pick, char_dd, [char_img, char_info])

    inputs = [msg, chat, model_dd, preset_dd, card_box, mem_box]
    send.click(respond, inputs, [chat, msg])
    msg.submit(respond, inputs, [chat, msg])
    clear.click(lambda: ([], ""), None, [chat, msg])
    start_btn.click(start_character, [char_dd, chat], [chat, card_box, char_info])

    def do_undo(history):
        return history[:-2] if len(history) >= 2 else []

    def do_retry(history, model_name, preset, card, memory_txt):
        if not history:
            return history
        # drop the last reply and re-ask the same question
        last_user = next((text_content(h.get("content")) for h in reversed(history)
                          if h["role"] == "user"), None)
        if last_user is None:
            return history
        trimmed = history[:-1] if history[-1]["role"] == "assistant" else history
        trimmed = trimmed[:-1] if trimmed and trimmed[-1]["role"] == "user" else trimmed
        out, _ = respond(last_user, trimmed, model_name, preset, card, memory_txt)
        return out

    undo.click(do_undo, chat, chat)
    retry.click(do_retry, [chat, model_dd, preset_dd, card_box, mem_box], chat)

    # A hidden, named Gradio endpoint for the PersonaMini web app.  Using a
    # normal Gradio event is important on ZeroGPU: Spaces owns the web server
    # and moves the decorated generate() call onto a GPU when needed.
    api_message = gr.Textbox(visible=False)
    api_history = gr.JSON(visible=False)
    api_model = gr.Textbox(visible=False)
    api_preset = gr.Textbox(visible=False)
    api_card = gr.Textbox(visible=False)
    api_memory = gr.Textbox(visible=False)
    api_key = gr.Textbox(visible=False)
    api_result = gr.JSON(visible=False)
    api_trigger = gr.Button(visible=False)

    def api_respond(message, history, model, preset, card, memory, supplied_key):
        secret = os.environ.get("PERSONAMINI_API_KEY", "")
        if not secret or not hmac.compare_digest(str(supplied_key or ""), secret):
            raise gr.Error("Invalid API key")
        model_name = next((name for name in MODELS if name.startswith(str(model or ""))), DEFAULT)
        history = history if isinstance(history, list) else []
        message = (message or "").strip()
        if not message:
            raise gr.Error("Message is required")

        pairs, user = [], None
        for item in history:
            if item.get("role") == "user":
                user = item.get("content", "")
            elif user is not None:
                pairs.append((user, item.get("content", "")))
                user = None
        pairs.append((message, ""))
        facts = [line.strip() for line in str(memory or "").splitlines() if line.strip()]
        facts += recall(pairs, str(card or ""))
        _m, tok, spec = get_model(model_name)
        max_new = max(200, spec["ctx"] // 2)
        reserve = min(max_new + 8, int(spec["ctx"] * 0.6))
        ids = build_prompt(str(card or ""), facts, pairs, spec["ctx"], reserve, tok)
        low = message.lower()
        n_lines = re.search(r"\b(\d+)\s*(lines?|words?|sentences?|items?)\b", low)
        min_new = 0 if (n_lines and int(n_lines.group(1)) <= 6) else (
            160 if any(word in low for word in LONG_WORDS) else 0)
        p = PRESETS.get(str(preset or "Default").title(), PRESETS["Default"])

        previous = ""
        for full_text in generate_stream(model_name, ids, p["temp"], p["top_k"], p["top_p"],
                                         p["rep"], p["no_repeat"], max_new, min_new):
            clean = full_text.strip()
            delta = clean[len(previous):] if clean.startswith(previous) else clean
            previous = clean
            if delta:
                yield {"delta": delta, "done": False}
        yield {"reply": previous or "…", "done": True}

    api_trigger.click(
        api_respond,
        [api_message, api_history, api_model, api_preset, api_card, api_memory, api_key],
        api_result,
        api_name="generate",
    )

    gr.Markdown(
        "---\nModels: "
        "[small](https://huggingface.co/bytestalkai/PersonaMini-1-small) · "
        "[medium](https://huggingface.co/bytestalkai/PersonaMini-1-medium) · "
        "[big](https://huggingface.co/bytestalkai/PersonaMini-1-big) — trained from scratch by "
        "BytesTalk. Replies are fiction from a very small model.")

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