| """Chat / text generation with the ONNX KnowLine model (raw onnxruntime, no torch). |
| |
| python generate.py --model ./model-cuda-int8 --cuda --prompt "Hello" # CUDA, int8 |
| python generate.py --model ./model-cuda-int4 --cuda --prompt "Hello" # CUDA, int4 |
| python generate.py --model ./model-cpu-int8 --prompt "Hello" # CPU, int8 |
| python generate.py --model ./model-cpu --interactive |
| python generate.py --model ./model-cpu --prompt "Hi" --max-new-tokens 100 --no-chat-template |
| |
| Uses onnx_raw.py (plain onnxruntime prefill + greedy decode). onnxruntime-genai |
| 0.19.0 cannot load this Qwen3.5 export on CPU, so it is intentionally not used. |
| """ |
|
|
| import argparse |
| import sys |
| import time |
| from pathlib import Path |
|
|
| HERE = Path(__file__).resolve().parent |
| sys.path.insert(0, str(HERE)) |
|
|
| from onnx_raw import OnnxLM |
|
|
|
|
| def eos_ids(model_dir): |
| import json |
|
|
| try: |
| cfg = json.loads((Path(model_dir) / "genai_config.json").read_text()) |
| eos = cfg["model"]["eos_token_id"] |
| return tuple(eos if isinstance(eos, list) else [eos]) |
| except (OSError, ValueError, KeyError): |
| return (248046, 248044) |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| p.add_argument("--model", default=str(HERE / "model-cpu"), help="build dir (model.onnx lives inside)") |
| p.add_argument("--onnx", default="model.onnx", help="onnx file name inside --model") |
| p.add_argument("--cuda", action="store_true", help="use CUDAExecutionProvider") |
| p.add_argument("--prompt", default=None) |
| p.add_argument("--interactive", action="store_true") |
| p.add_argument("--max-new-tokens", type=int, default=512) |
| p.add_argument("--no-chat-template", action="store_true", help="use the prompt verbatim (no chat template)") |
| p.add_argument("--threads", type=int, default=0, help="ORT intra-op threads (0 = default)") |
| a = p.parse_args() |
|
|
| from transformers import AutoTokenizer |
|
|
| tok = AutoTokenizer.from_pretrained(str(HERE / "orig_tokenizer"), local_files_only=True) |
| providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if a.cuda else ["CPUExecutionProvider"] |
| print(f"Loading {a.model}/{a.onnx} {providers[0]} ...", flush=True) |
| lm = OnnxLM(str(Path(a.model) / a.onnx), providers=providers, intra_threads=a.threads) |
| eos = eos_ids(a.model) |
|
|
| def run(prompt): |
| if a.no_chat_template: |
| ids = tok.encode(prompt, add_special_tokens=False) |
| else: |
| text = tok.apply_chat_template([{"role": "user", "content": prompt}], |
| tokenize=False, add_generation_prompt=True, enable_thinking=False) |
| ids = tok.encode(text, add_special_tokens=False) |
| print(f" prompt tokens: {len(ids)}", flush=True) |
| t0 = time.perf_counter() |
| n_first = {"done": False} |
| t_first = {} |
|
|
| def on_token(t): |
| if not n_first["done"]: |
| n_first["done"] = True |
| t_first["t"] = time.perf_counter() - t0 |
| sys.stdout.write(tok.decode([t])) |
| sys.stdout.flush() |
|
|
| out = lm.generate(ids, a.max_new_tokens, eos, on_token=on_token) |
| total = time.perf_counter() - t0 |
| print() |
| n = len(out) |
| print(f" tokens={n} TTFT={t_first.get('t', total) * 1000:.0f}ms " |
| f"TPS={max(n - 1, 1) / max(total - t_first.get('t', 0), 1e-9):.1f} total={total:.1f}s") |
|
|
| if a.interactive: |
| print("Interactive mode ('quit' to stop).") |
| while True: |
| try: |
| prompt = input("You: ").strip() |
| except EOFError: |
| break |
| if prompt.lower() in ("quit", "exit") or not prompt: |
| break |
| run(prompt) |
| print("-" * 50) |
| elif a.prompt: |
| run(a.prompt) |
| else: |
| p.print_help() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
