| """The tiny-agent backend: moth-112m's own architecture (tiny_agent/model.py: CSA attention, Engram tables), the |
| default. Loading is exactly what moth.serve.Moth always did (fp32 + memory-mapped bf16 Engram tables + int8 Linear |
| weights on the CPU, bf16 on a GPU, one KV cache for the session). |
| |
| It is the one stateful backend: moth.serve.Moth keeps the KV cache across turns and feeds only each new message |
| (Moth._feed / _prefill_all / _turn, tiny_agent.generate.Decoder with captured device graphs on CUDA/XPU and the |
| phrase ban as a logits hook), so a turn costs only its new tokens. That path stays in Moth itself, untouched; this |
| class owns the loading and the capability flags. generate() is the stateless form of the same model (fresh cache, |
| whole prompt) for callers that just want a reply to a message list. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| from tokenizers import Tokenizer |
|
|
| from moth.backends.base import Backend |
|
|
|
|
| class TinyAgentBackend(Backend): |
| name = "tiny-agent" |
| stateful = True |
| tools = True |
| graphs = True |
| phrase_ban = True |
|
|
| def __init__(self, a): |
| from moth import serve |
| from tiny_agent.checkpoint import load_model |
| from tiny_agent.generate import KVCache |
| self.a = a |
| self.tok = Tokenizer.from_file(a.tokenizer) |
| dtype = torch.float32 if a.device == "cpu" else torch.bfloat16 |
| |
| self.model = load_model(a.model, device=a.device, dtype=dtype, low_mem=a.device == "cpu").eval() |
| quant = self.quant = serve.resolve_quant(a) |
| if quant != "none": |
| if a.device != "cpu": |
| raise SystemExit("--quant is for --device cpu") |
| from tiny_agent.generate import quantize_cpu |
| quantize_cpu(self.model, quant) |
| self.kv_dtype = dtype |
| |
| self.kv = KVCache(self.model, 1, serve.MAX_LEN, a.device, dtype=dtype) |
|
|
| def tokenizer(self): |
| return self.tok |
|
|
| def _ids(self, messages): |
| from tiny_agent.chat import render |
| return self.tok.encode(render(messages, add_generation_prompt=True), add_special_tokens=False).ids |
|
|
| def count(self, messages): |
| return len(self._ids(messages)) |
|
|
| @torch.no_grad() |
| def generate(self, messages, *, max_tokens, temperature, top_p, rep_penalty=1.0, pen_ids=(), ban_tools=False, |
| logits_hook=None, stop=()): |
| from moth import serve |
| from tiny_agent.generate import Decoder, KVCache, append |
| ids = self._ids(messages) |
| cache = KVCache(self.model, 1, serve.MAX_LEN, self.a.device, dtype=self.kv_dtype) |
| logits = append(self.model, cache, torch.tensor([ids], device=self.a.device), [len(ids)]) |
| ban = [self.tok.token_to_id(t) for t in ("<tool_call>", "<think>")] if ban_tools else () |
| out = Decoder(self.model, cache, graph=False).generate( |
| logits, max_tokens, self.tok.token_to_id("<|im_end|>"), temperature=temperature, top_p=top_p, |
| pen=rep_penalty, pen_ids=set(pen_ids), ban_ids=ban, logits_hook=logits_hook) |
| from moth.backends.base import cut_at_stops |
| return cut_at_stops(self.tok.decode(out, skip_special_tokens=False), stop) |
|
|
