"""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())