Spaces:
Running on Zero
Running on Zero
| """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 | |
| 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) | |
| 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()) | |