| """Chat with Ling-mini-2.0 in a terminal, on the megakernel. |
| |
| python chat.py [--max-new 1024] |
| |
| Greedy decoding, thinking off. Each answer re-reads only the new part of the conversation (the rest |
| is still in the kernel's cache). Windows lowers an idle GPU's clocks, so the kernel warms up at start |
| and before each answer. The speed shown is the kernel's own while it writes. An empty line or |
| "exit" quits. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| import time |
|
|
| from ling_kernel import HERE, LingKernel |
|
|
| CHUNK = 32 |
|
|
|
|
| def calibrate(mk: LingKernel, tok, budget: float = 20.0) -> float: |
| ids = tok("Hello")["input_ids"] |
| first = mk.forced(ids)[-1] |
| start, best, since = time.perf_counter(), float("inf"), time.perf_counter() |
| while time.perf_counter() - start < budget: |
| s = mk._launch([first], len(ids), CHUNK, False) |
| if s < 0.97 * best: |
| best, since = s, time.perf_counter() |
| elif time.perf_counter() - since > 1.0: |
| break |
| mk.cached = [] |
| return best |
|
|
|
|
| def warm_up(mk: LingKernel, first: int, pos: int, budget: float = 3.0) -> None: |
| start = time.perf_counter() |
| while time.perf_counter() - start < budget: |
| if mk._launch([first], pos, CHUNK, False) < 1.15 * mk.best_launch: |
| break |
|
|
|
|
| def answer(mk: LingKernel, tok, eos: set[int], messages: list[dict], max_new: int) -> str: |
| ids = tok.apply_chat_template(messages, tokenize=True, add_generation_prompt=True) |
| if hasattr(ids, "keys"): |
| ids = ids["input_ids"] |
| ids = list(ids) |
| if len(ids) + max_new + CHUNK > mk.max_pos: |
| print(f"[the conversation is too long for the {mk.max_pos}-token cache; type 'exit' and start again]") |
| return "" |
| n = 0 |
| while n < min(len(ids) - 1, len(mk.cached)) and ids[n] == mk.cached[n]: |
| n += 1 |
| preds = mk.forced(ids[n:], pos0=n) |
| warm_up(mk, preds[-1], len(ids)) |
| out, seconds, steps, pos, shown = [preds[-1]], 0.0, 0, len(ids), "" |
| print("Ling> ", end="", flush=True) |
| while out[-1] not in eos and len(out) < max_new: |
| k = min(CHUNK, max_new - len(out)) |
| seconds += mk._launch([out[-1]], pos, k, False) |
| steps += k |
| pos += k |
| for t in mk.state.out[:k].tolist(): |
| out.append(t) |
| if t in eos: |
| break |
| text = tok.decode([t for t in out if t not in eos], skip_special_tokens=True) |
| print(text[len(shown):], end="", flush=True) |
| shown = text |
| mk.cached = ids + out[:-1] |
| speed = f"{steps / seconds:.0f} tokens/s" if seconds else "-" |
| print(f"\n[{len(out)} tokens; the kernel writes {speed}; read {len(ids) - n} new prompt tokens]\n", flush=True) |
| return shown.strip() |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--max-new", type=int, default=1024) |
| args = ap.parse_args() |
| sys.stdout.reconfigure(encoding="utf-8") |
| from transformers import AutoTokenizer |
| tok = AutoTokenizer.from_pretrained(HERE) |
| print("loading the Ling-mini-2.0 megakernel (4-bit) ...", flush=True) |
| mk = LingKernel(max_pos=4096, log=lambda s: print(s, flush=True)) |
| eos = {i for i in (tok.convert_tokens_to_ids("<|role_end|>"), tok.eos_token_id, mk.config.get("eos_token_id")) |
| if isinstance(i, int) and i >= 0} |
| print("warming up the GPU ...", flush=True) |
| mk.best_launch = calibrate(mk, tok) |
| print(f" {CHUNK / mk.best_launch:.0f} tokens/s at full speed", flush=True) |
| history: list[dict] = [] |
| print("type a message; an empty line or 'exit' quits\n", flush=True) |
| while True: |
| try: |
| text = input("you> ").strip() |
| except EOFError: |
| break |
| if not text or text.lower() in ("exit", "quit"): |
| break |
| messages = history + [{"role": "user", "content": text}] |
| reply = answer(mk, tok, eos, messages, args.max_new) |
| history = messages + [{"role": "assistant", "content": reply}] |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
