| """The transformers backend: any causal LM `transformers` can load (a local folder or an HF repo id), on the CPU or a |
| GPU, behind moth. Optional extra: pip install "moth-companion[transformers]" (transformers is imported only here). |
| |
| The prompt is the tokenizer's own chat template (chat_format "template", the default) or moth's ChatML |
| (chat_format "chatml", for models trained on it). Sampling is moth's own (tiny_agent.generate.sample_top_p: the |
| repetition penalty on moth's recent words, temperature, top-p), one token at a time with the model's KV cache, so |
| moth's phrase ban (moth/norepeat.py) works as a logits hook here too. Stateless across turns: each turn re-reads the |
| chat (small models: milliseconds to a second on the CPU). |
| """ |
| from __future__ import annotations |
|
|
| import torch |
|
|
| from moth.backends.base import Backend, _Enc, chat_messages, cut_at_stops |
|
|
| STOP_TOKENS = ("<|im_end|>", "<|eot_id|>", "<end_of_turn>", "<|end|>", "<|endoftext|>", "</s>") |
| KEEP_TEXT = ("<think>", "</think>", "<tool_call>", "</tool_call>") |
|
|
|
|
| class HFTokenizer: |
| """A transformers tokenizer with the `tokenizers.Tokenizer` methods moth.serve uses. SentencePiece's word-start |
| "▁" reads as byte-level BPE's "Ġ" in get_vocab(), so the phrase ban's word classes (norepeat.vocab_classes) find |
| word starts in both kinds of vocabulary.""" |
|
|
| def __init__(self, tok): |
| self.tok = tok |
| self._vocab = None |
|
|
| def encode(self, text, add_special_tokens=False): |
| return _Enc(self.tok(text, add_special_tokens=add_special_tokens)["input_ids"]) |
|
|
| def decode(self, ids, skip_special_tokens=False): |
| return self.tok.decode(list(ids), skip_special_tokens=skip_special_tokens) |
|
|
| def token_to_id(self, s): |
| return self.tok.get_vocab().get(s) |
|
|
| def get_vocab(self): |
| if self._vocab is None: |
| self._vocab = {("Ġ" + s[1:] if s.startswith("▁") else s): i for s, i in self.tok.get_vocab().items()} |
| return self._vocab |
|
|
| def get_added_tokens_decoder(self): |
| return dict(getattr(self.tok, "added_tokens_decoder", {}) or {}) |
|
|
|
|
| class TransformersBackend(Backend): |
| name = "transformers" |
| tools = True |
| phrase_ban = True |
|
|
| def __init__(self, m, a): |
| try: |
| import transformers |
| except ImportError: |
| raise SystemExit('moth: this model needs the transformers backend: pip install "moth-companion[transformers]"' |
| " (or pip install transformers)") |
| self.m = m |
| self.device = a.device |
| src = m.model or m.folder |
| if not src: |
| raise SystemExit("moth: the transformers backend needs a model (a folder or an HF repo id) in moth.json") |
| self.tok = transformers.AutoTokenizer.from_pretrained(src) |
| dtype = torch.float32 if a.device == "cpu" else torch.bfloat16 |
| try: |
| model = transformers.AutoModelForCausalLM.from_pretrained(src, dtype=dtype) |
| except TypeError: |
| model = transformers.AutoModelForCausalLM.from_pretrained(src, torch_dtype=dtype) |
| self.model = model.to(a.device).eval() |
| if m.chat_format == "template" and not getattr(self.tok, "chat_template", None): |
| raise SystemExit(f"moth: {src} has no chat template; set \"chat_format\": \"chatml\" in its moth.json " |
| "(or pick a chat/instruct model)") |
| vocab = self.tok.get_vocab() |
| self.stop_ids = {i for i in (self.tok.eos_token_id,) if isinstance(i, int)} |
| gen = getattr(self.model, "generation_config", None) |
| eos = getattr(gen, "eos_token_id", None) |
| self.stop_ids |= set(eos if isinstance(eos, (list, tuple)) else [eos] if isinstance(eos, int) else []) |
| self.stop_ids |= {vocab[t] for t in STOP_TOKENS if t in vocab} |
| self.tool_ids = [vocab[t] for t in ("<tool_call>", "<think>") if t in vocab] |
| self.drop = [t for t in (self.tok.all_special_tokens or []) if t not in KEEP_TEXT] |
| self.adapter = HFTokenizer(self.tok) |
|
|
| def tokenizer(self): |
| return self.adapter |
|
|
| def prompt(self, messages) -> str: |
| if self.m.chat_format == "chatml": |
| from tiny_agent.chat import render |
| return render(messages, add_generation_prompt=True) |
| return self.tok.apply_chat_template(chat_messages(messages), tokenize=False, add_generation_prompt=True) |
|
|
| def _ids(self, messages): |
| return self.tok(self.prompt(messages), add_special_tokens=False)["input_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 tiny_agent.generate import sample_top_p |
| ids = self._ids(messages) |
| out = self.model(input_ids=torch.tensor([ids], device=self.device), use_cache=True) |
| past, logits = out.past_key_values, out.logits[:, -1, :] |
| V = logits.size(-1) |
| pen_list = [i for i in pen_ids if 0 <= i < V] |
| pm = None |
| if pen_list and rep_penalty != 1.0: |
| pm = torch.zeros(V, dtype=torch.bool, device=logits.device) |
| pm[pen_list] = True |
| bm = None |
| if ban_tools and self.tool_ids: |
| bm = torch.zeros(V, dtype=torch.bool, device=logits.device) |
| bm[self.tool_ids] = True |
| gen: list[int] = [] |
| for _ in range(max_tokens): |
| if logits_hook is not None: |
| logits = logits_hook(gen, logits) |
| t = int(sample_top_p(logits, temperature, top_p, rep_penalty, pm, bm)) |
| if t in self.stop_ids: |
| break |
| gen.append(t) |
| if stop and any(s in self.tok.decode(gen[-16:]) for s in stop): |
| break |
| out = self.model(input_ids=torch.tensor([[t]], device=self.device), past_key_values=past, use_cache=True) |
| past, logits = out.past_key_values, out.logits[:, -1, :] |
| text = self.tok.decode(gen, skip_special_tokens=False) |
| for s in self.drop: |
| text = text.replace(s, "") |
| return cut_at_stops(text, stop).strip() |
|
|
