File size: 5,233 Bytes
5cf6791 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """The backend interface moth.serve talks to, and helpers the general backends share. Stdlib only.
A backend loads a model and turns moth's conversation (tiny_agent.chat's canonical messages: system / user /
assistant with tool_calls / tool with results) into the model's next assistant turn:
tokenizer() an object with the bits of a `tokenizers.Tokenizer` moth.serve uses: encode(text,
add_special_tokens=False).ids, decode(ids, skip_special_tokens=False), token_to_id(s), get_vocab(),
get_added_tokens_decoder(). It drives the context budget, the repetition penalty's word ids and the
phrase ban's word classes (moth/norepeat.py).
count(messages) tokens the whole prompt takes (the context budget; approximate for openai)
generate(messages, max_tokens, temperature, top_p, rep_penalty, pen_ids, ban_tools, logits_hook, stop)
the assistant turn as text (may hold <tool_call> JSON when the manifest enables tools).
pen_ids: token ids the repetition penalty hits; logits_hook: moth's phrase ban
(norepeat.BanHook(out_ids, logits) -> logits), only for backends with phrase_ban.
Capabilities: stateful (moth.serve keeps a KV cache across turns and feeds only the new text: tiny-agent),
tools (can follow moth's <tool_call> format; the manifest still has to switch tools on), graphs (device graphs),
phrase_ban (a logits hook works: needs token-level sampling).
"""
from __future__ import annotations
import json
import re
class Backend:
name = "?"
stateful = False
tools = False
graphs = False
phrase_ban = False
def tokenizer(self):
raise NotImplementedError
def count(self, messages: list[dict]) -> int:
raise NotImplementedError
def generate(self, messages: list[dict], *, max_tokens: int, temperature: float, top_p: float,
rep_penalty: float = 1.0, pen_ids=(), ban_tools: bool = False, logits_hook=None, stop=()) -> str:
raise NotImplementedError
def capabilities(self) -> dict:
return {k: getattr(self, k) for k in ("stateful", "tools", "graphs", "phrase_ban")}
def chat_messages(messages: list[dict]) -> list[dict]:
"""moth's canonical messages as plain {"role", "content"} chat messages for a chat template / an OpenAI server.
Tool calls stay in the assistant text as moth writes them (<tool_call> JSON); tool results go back as a user
message of <tool_response> blocks (OpenAI's role "tool" would need a tool_call_id, and many small-model templates
don't know it). Moth's own bookkeeping keys ("session") never reach the model."""
out = []
for m in messages:
r = m["role"]
if r in ("system", "user"):
out.append({"role": r, "content": m["content"]})
elif r == "assistant":
s = m.get("content") or ""
for c in m.get("tool_calls") or []:
if "name" in c:
s += ("\n" if s else "") + "<tool_call>\n" + json.dumps(
{"name": c["name"], "arguments": c.get("arguments", {})}, ensure_ascii=False) + "\n</tool_call>"
out.append({"role": "assistant", "content": s})
elif r == "tool":
out.append({"role": "user", "content": "".join(f"<tool_response>\n{x}\n</tool_response>"
for x in m["results"])})
return out
def cut_at_stops(text: str, stop=()) -> str:
for s in stop:
i = text.find(s)
if i >= 0:
text = text[:i]
return text
class _Enc:
__slots__ = ("ids",)
def __init__(self, ids):
self.ids = ids
class PseudoTokenizer:
"""For backends without a local tokenizer (openai): words and punctuation get stable ids for this process, so the
context budget, the repetition-penalty bookkeeping and tool-result cutting work. No vocabulary (no phrase ban)."""
WORD = re.compile(r"\w+|[^\w\s]", re.U)
MAX = 200_000
def __init__(self):
self.ids: dict[str, int] = {}
self.words: list[str] = []
def encode(self, text: str, add_special_tokens: bool = False) -> _Enc:
out = []
for w in self.WORD.findall(text):
i = self.ids.get(w)
if i is None:
if len(self.words) >= self.MAX: # a server that runs for weeks: start over, ids are per turn anyway
self.ids.clear()
self.words.clear()
i = self.ids[w] = len(self.words)
self.words.append(w)
out.append(i)
return _Enc(out)
def decode(self, ids, skip_special_tokens: bool = False) -> str:
return " ".join(self.words[i] for i in ids if 0 <= i < len(self.words))
def token_to_id(self, s: str):
return None
def get_vocab(self) -> dict:
return {}
def get_added_tokens_decoder(self) -> dict:
return {}
def approx_tokens(messages: list[dict]) -> int:
"""~3 characters per token (a little high for English, so the budget errs on the safe side) + 4 per message."""
return sum(len(m.get("content") or "") for m in chat_messages(messages)) // 3 + 4 * len(messages) + 3
|