mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 09:36:08 +00:00
215 lines
13 KiB
Python
215 lines
13 KiB
Python
from __future__ import annotations
|
|
import sys, argparse, codecs, itertools, typing, re, unicodedata, json, time
|
|
from typing import TYPE_CHECKING
|
|
from tinygrad import nn
|
|
from tinygrad.uop.ops import UOp, Ops
|
|
from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, Context, fetch, profile_marker, getenv
|
|
from tinygrad.llm.model import Transformer
|
|
if TYPE_CHECKING:
|
|
import jinja2
|
|
|
|
class SimpleTokenizer:
|
|
def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int], preset:str="llama3",
|
|
bos_id:int|None=None, eos_id:int=0, eot_id:int|None=None):
|
|
preset = {"qwen35":"qwen2","qwen35moe":"qwen2"}.get(preset, preset)
|
|
if preset not in ("llama3","llama-v3","llama-bpe","qwen2","olmo","kimi-k2","tekken","glm4","gpt-4o"):
|
|
raise ValueError(f"Invalid tokenizer preset '{preset}'")
|
|
# https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
|
|
bs = [*range(33, 127), *range(161, 173), *range(174, 256)] # bytes that map to themselves
|
|
self._byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
|
|
|
|
# https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
|
|
# each limit is one past the category's max codepoint (Z→U+3000, N→U+1FBF9, L→U+323AF, M→U+E01EF)
|
|
# compact adjacent codepoints into ranges: listing them all makes re spend seconds on large prompts
|
|
def ucat_range(pre:str|tuple[str, ...]) -> str:
|
|
limits = {"Z": 0x3001, "N": 0x1fbfa, "L": 0x323b0, "M": 0xe01f0}
|
|
limit = max(limits[p if p in limits else p[0]] for p in (pre if isinstance(pre, tuple) else (pre,)))
|
|
cps = enumerate(cp for cp in range(limit) if unicodedata.category(chr(cp)).startswith(pre))
|
|
runs = [list(g) for _, g in itertools.groupby(cps, lambda e: e[1]-e[0])]
|
|
return "".join(re.escape(chr(g[0][1])) + (f"-{re.escape(chr(g[-1][1]))}" if len(g) > 1 else "") for g in runs)
|
|
r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
|
|
contr, r_l, r_n = "(?i:'s|'t|'re|'ve|'m|'ll|'d)", f"[^\\r\\n{r_p_N}{r_p_L}]?", f"[{r_p_N}]" if preset == "tekken" else f"[{r_p_N}]{{1,3}}"
|
|
r_p, r_w, r_t = f" ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*", f"{contr}|{r_l}[{r_p_L}]+", f"[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+"
|
|
if preset in ("tekken", "gpt-4o"):
|
|
r_up, r_lo = ucat_range(("Lu","Lt","Lm","Lo","M")), ucat_range(("Ll","Lm","Lo","M"))
|
|
sfx = f"{contr}?" if preset == "gpt-4o" else ""
|
|
r_p, r_w = f" ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n/]*", f"{r_l}[{r_up}]*[{r_lo}]+{sfx}|{r_l}[{r_up}]+[{r_lo}]*{sfx}"
|
|
self._split_to_word = re.compile(f"{r_w}|{r_n}|{r_p}|{r_t}")
|
|
self._split_to_sentence = re.compile("|".join(re.escape(tok) for tok in special_tokens.keys()) if special_tokens else r"(?!)")
|
|
|
|
self._normal_tokens = {bytes(self._byte_decoder[c] for c in tok): tid for tok, tid in normal_tokens.items()}
|
|
self._special_tokens = special_tokens
|
|
self._tok2bytes = {tid: tok for tok, tid in self._normal_tokens.items()} | {tid: tok.encode() for tok, tid in self._special_tokens.items()}
|
|
self.preset = preset
|
|
self.bos_id, self.eos_id, self.eot_id = bos_id, eos_id, eot_id
|
|
|
|
@staticmethod
|
|
def from_gguf_kv(kv:dict):
|
|
# https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L1818-L1820
|
|
vocab: typing.Iterable[tuple[str, int]] = ((tok, idx) for idx, tok in enumerate(kv["tokenizer.ggml.tokens"]))
|
|
normal_tokens, special_tokens = partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1)
|
|
special_tokens_dict = dict(special_tokens)
|
|
return SimpleTokenizer(dict(normal_tokens), special_tokens_dict, kv["tokenizer.ggml.pre"],
|
|
bos_id=kv.get('tokenizer.ggml.bos_token_id') if kv.get('tokenizer.ggml.add_bos_token', True) else None,
|
|
eos_id=kv.get('tokenizer.ggml.eos_token_id', 0), eot_id=kv.get('tokenizer.ggml.eot_token_id', special_tokens_dict.get('<|im_end|>')))
|
|
|
|
def _encode_word(self, word:bytes) -> list[int]:
|
|
if (early_token:=self._normal_tokens.get(word)) is not None: return [early_token]
|
|
parts = [bytes([b]) for b in word]
|
|
# greedily merge any parts that we can
|
|
while True:
|
|
i = min([(sys.maxsize, -1)] + [(self._normal_tokens.get(parts[j]+parts[j+1], sys.maxsize), j) for j in range(len(parts)-1)])[1]
|
|
if i == -1: break
|
|
parts[i:i+2] = [parts[i] + parts[i+1]]
|
|
try: return [self._normal_tokens[p] for p in parts]
|
|
except KeyError: raise RuntimeError("token not found")
|
|
def _encode_sentence(self, chunk:str) -> list[int]:
|
|
return [tok for word in self._split_to_word.findall(chunk) for tok in self._encode_word(word.encode())]
|
|
def encode(self, text:str) -> list[int]:
|
|
tokens: list[int] = []
|
|
pos = 0
|
|
for match in self._split_to_sentence.finditer(text):
|
|
tokens.extend(self._encode_sentence(text[pos:match.start(0)]) + [self._special_tokens[text[match.start(0):match.end(0)]]])
|
|
pos = match.end(0)
|
|
return tokens + self._encode_sentence(text[pos:])
|
|
|
|
def decode(self, ids:list[int]) -> str: return b''.join(self._tok2bytes[tid] for tid in ids).decode(errors='replace')
|
|
def stream_decoder(self) -> typing.Callable[..., str]:
|
|
dec = codecs.getincrementaldecoder('utf-8')('replace')
|
|
def _decode(tid:int|None=None) -> str: return dec.decode(self._tok2bytes[tid]) if tid is not None else dec.decode(b'', final=True)
|
|
return _decode
|
|
def is_end(self, token_id:int) -> bool: return token_id in (self.eos_id, self.eot_id)
|
|
|
|
models = {
|
|
"llama3.2:1b": "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q6_K.gguf",
|
|
"llama3.2:1b-q4": "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q4_K_M.gguf",
|
|
"llama3.2:3b": "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-Q6_K.gguf",
|
|
"llama3.2:3b-f16": "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-f16.gguf",
|
|
"llama3.1:8b": "https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF/resolve/main/Meta-Llama-3.1-8B-Instruct-Q8_0.gguf",
|
|
"qwen3:0.6b": "https://huggingface.co/Qwen/Qwen3-0.6B-GGUF/resolve/main/Qwen3-0.6B-Q8_0.gguf",
|
|
"qwen3:1.7b": "https://huggingface.co/unsloth/Qwen3-1.7B-GGUF/resolve/main/Qwen3-1.7B-Q4_K_M.gguf",
|
|
"qwen3:8b": "https://huggingface.co/Qwen/Qwen3-8B-GGUF/resolve/main/Qwen3-8B-Q4_K_M.gguf",
|
|
"qwen3:30b-a3b": "https://huggingface.co/Qwen/Qwen3-30B-A3B-GGUF/resolve/main/Qwen3-30B-A3B-Q4_K_M.gguf",
|
|
"qwen3.5:0.8b": "https://huggingface.co/unsloth/Qwen3.5-0.8B-GGUF/resolve/main/Qwen3.5-0.8B-Q8_0.gguf",
|
|
"qwen3.5:4b": "https://huggingface.co/unsloth/Qwen3.5-4B-GGUF/resolve/main/Qwen3.5-4B-Q4_K_M.gguf",
|
|
"qwen3.5:9b": "https://huggingface.co/unsloth/Qwen3.5-9B-GGUF/resolve/main/Qwen3.5-9B-Q4_K_M.gguf",
|
|
"qwen3.6:27b": "https://huggingface.co/unsloth/Qwen3.6-27B-GGUF/resolve/main/Qwen3.6-27B-Q4_K_M.gguf",
|
|
"qwen3.6:35b-a3b": "https://huggingface.co/unsloth/Qwen3.6-35B-A3B-GGUF/resolve/main/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf",
|
|
# pinned to the last revision with the plain IQ4_XS quant: the UD replacement uses Q3_K tensors the loader doesn't support
|
|
"qwen3.8:27b": "https://huggingface.co/unsloth/Qwen3.8-27B-GGUF/resolve/b62a80264f8b0c1bb849ee1c9c487415ebeca194/Qwen3.8-27B-IQ4_XS.gguf",
|
|
"olmoe": "https://huggingface.co/allenai/OLMoE-1B-7B-0924-Instruct-GGUF/resolve/main/olmoe-1b-7b-0924-instruct-q4_k_m.gguf",
|
|
"moonlight": "https://huggingface.co/gabriellarson/Moonlight-16B-A3B-Instruct-GGUF/resolve/main/Moonlight-16B-A3B-Instruct-Q4_K_M.gguf",
|
|
"glm-4.7-flash": "https://huggingface.co/unsloth/GLM-4.7-Flash-GGUF/resolve/main/GLM-4.7-Flash-Q4_K_M.gguf",
|
|
}
|
|
|
|
class FallbackTemplate:
|
|
# minimal jinja2.Template-compatible chat template without jinja2, no tool calling support
|
|
def __init__(self, tok:SimpleTokenizer): self.tok = tok
|
|
def role(self, role:str) -> str:
|
|
if self.tok.preset == 'olmo': return "<|" + role + "|>\n" # OLMoE Instruct format
|
|
if self.tok.preset == 'kimi-k2': return "<|im_" + role + "|>" + role + "<|im_middle|>"
|
|
if self.tok.preset == 'qwen2': return "<|im_start|>" + role + "\n"
|
|
if self.tok.preset == 'glm4': return "<|" + role + "|>"
|
|
if self.tok.preset == 'tekken':
|
|
if role == 'user': return "[INST]"
|
|
if role == 'assistant': return ""
|
|
raise ValueError(f"Unsupported role '{role}' for tokenizer preset '{self.tok.preset}'")
|
|
return "<|start_header_id|>" + role + "<|end_header_id|>\n\n"
|
|
def end_turn(self) -> str:
|
|
if self.tok.preset == 'olmo': return "\n"
|
|
if self.tok.preset == 'kimi-k2': return self.tok.decode([self.tok.eos_id])
|
|
if self.tok.preset == 'qwen2': return self.tok.decode([self.tok.eos_id]) + "\n"
|
|
if self.tok.preset == 'glm4': return ""
|
|
if self.tok.preset == 'tekken': return "[/INST]"
|
|
return self.tok.decode([self.tok.eos_id])
|
|
def render(self, messages:list[dict], tools=None, add_generation_prompt:bool=True, preserve_thinking:bool=False) -> str:
|
|
out = self.tok.decode([] if self.tok.bos_id is None else [self.tok.bos_id]) + ("<sop>" if self.tok.preset == 'glm4' else "")
|
|
for msg in messages:
|
|
out += self.role(msg["role"])
|
|
content = msg.get("content")
|
|
if isinstance(content, str): out += content
|
|
elif isinstance(content, list):
|
|
for c in content:
|
|
if c["type"] == "text": out += c["text"]
|
|
else: raise RuntimeError(f"unhandled type: {c['type']}")
|
|
elif content is not None: raise RuntimeError(f"unknown content type: {type(content)}")
|
|
out += self.end_turn()
|
|
return out + self.role("assistant") if add_generation_prompt else out
|
|
|
|
from tinygrad.llm.serve import LLMServer
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model", "-m", default=list(models.keys())[0], help=f"Model choice ({', '.join(models.keys())}) or path to a local GGUF file")
|
|
parser.add_argument("--max_context", type=int, default=4096, help="Max Context Length")
|
|
parser.add_argument("--serve", nargs='?', type=int, const=8000, metavar="PORT", help="Run OpenAI compatible API (optional port, default 8000)")
|
|
parser.add_argument("--warmup", action="store_true", help="warmup the JIT")
|
|
parser.add_argument("--benchmark", nargs='?', type=int, const=20, metavar="COUNT", help="Benchmark tok/s (optional count, default 20)")
|
|
parser.add_argument("--no_chat_template", action="store_true", help="Don't use the model's chat template, always use the fallback template")
|
|
args = parser.parse_args()
|
|
|
|
# load the model
|
|
with Context(DEBUG=max(DEBUG.value, 2 if args.serve else 0)):
|
|
model, kv = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
|
|
model_name = kv.get('general.name') or kv.get('general.basename') or args.model
|
|
file_sizes = [y.nbytes() for y in UOp.sink(*[x.uop for x in nn.state.get_parameters(model)]).toposort() if y.op is Ops.BUFFER]
|
|
print(f"using model \"{model_name}\" with {sum(file_sizes):,} bytes and {sum(x.numel() for x in nn.state.get_parameters(model)):,} params, "
|
|
f"max context {args.max_context} on {nn.state.get_parameters(model)[0].device}")
|
|
|
|
# get tokenizer
|
|
tok = SimpleTokenizer.from_gguf_kv(kv)
|
|
|
|
# use the model's chat template if jinja2 is available (enables model-specific formatting)
|
|
template: jinja2.Template|FallbackTemplate = FallbackTemplate(tok)
|
|
if not args.no_chat_template and (ct := kv.get('tokenizer.chat_template')) is not None:
|
|
try:
|
|
import jinja2
|
|
env = jinja2.Environment()
|
|
env.filters['tojson'] = lambda obj, **kwargs: json.dumps(obj, **kwargs) # jinja2's tojson escapes <>& for HTML safety
|
|
env.globals['raise_exception'] = lambda msg: (_ for _ in ()).throw(RuntimeError(msg))
|
|
env.globals['strftime_now'] = lambda fmt: time.strftime(fmt)
|
|
env.globals['bos_token'] = tok.decode([tok.bos_id]) if tok.bos_id is not None else ""
|
|
env.globals['eos_token'] = tok.decode([tok.eos_id])
|
|
template = env.from_string(ct)
|
|
except ImportError: print("warning: jinja2 is not installed, the model's chat template is disabled")
|
|
|
|
# warmup the JIT
|
|
if args.warmup or args.serve:
|
|
with Context(DEBUG=max(DEBUG.value, 1)): model.warmup()
|
|
|
|
# start server
|
|
if args.serve: LLMServer(('', args.serve), model, model_name, tok, template).serve_forever()
|
|
|
|
# do benchmark
|
|
if args.benchmark is not None:
|
|
gen = model.generate(toks:=[tok.bos_id or 0])
|
|
for i in range(args.benchmark):
|
|
profile_marker(f"decode @ {i}")
|
|
GlobalCounters.reset()
|
|
if (log:=getenv("BENCHMARK_LOG", "")): from extra.bench_log import WallTimeEvent, BenchEvent
|
|
with Timing(on_exit=lambda x: f", {1e9/x:6.2f} tok/s, {GlobalCounters.global_mem/x:7.2f} GB/s,"
|
|
f" {GlobalCounters.global_mem//1000000}/{GlobalCounters.mem_used//1000000} MB -- "+\
|
|
tok.decode(toks).replace("\n", "\\n")):
|
|
if log:
|
|
with WallTimeEvent(BenchEvent.STEP): next(gen)
|
|
else: next(gen)
|
|
exit(0)
|
|
|
|
# interactive chat
|
|
messages: list[dict] = []
|
|
while 1:
|
|
try: messages.append({"role":"user", "content":input('>>> ')})
|
|
except EOFError: break
|
|
ids = tok.encode(template.render(messages=messages, add_generation_prompt=True))
|
|
reply, dec = "", tok.stream_decoder()
|
|
for next_id in model.generate(ids):
|
|
if tok.is_end(next_id):
|
|
sys.stdout.write(dec() + "\n\n")
|
|
break
|
|
reply += (piece := dec(next_id))
|
|
sys.stdout.write(piece)
|
|
sys.stdout.flush()
|
|
messages.append({"role":"assistant", "content":reply})
|
|
|
|
if __name__ == "__main__": main()
|