From 7ec9ced71b67729e097eb3033fbfb5c03d895c7d Mon Sep 17 00:00:00 2001 From: bong-water-water-bong <277547417+bong-water-water-bong@users.noreply.github.com> Date: Tue, 22 Sep 2026 17:21:54 -0300 Subject: [PATCH] Server: 1bit-server speaks Lemonade's backend protocol - engine/server/main.cpp: 1bit-server. It takes Lemonade's launch flags (-m --ctx-size --device --port --jinja --metrics --parallel 1) and serves: /health (503 while loading), /v1/models, /v1/chat/completions and /v1/completions, streaming (SSE, [DONE]) and not. Responses carry OpenAI usage and llama-server timings. - generator: prefix-cached generation loop (cache_n), end-of-turn tokens. - text_stream: UTF-8-safe deltas, stop strings across tokens, routed to reasoning_content. - chat_template: the GGUF's Jinja template via minja. - openai: request parsing (rejects n>1, non-text parts, over-long prompts) and response shapes. - core/sampler: temperature/top-k/top-p/min-p/seed, llama.cpp chain order. - CpuModel: configurable n_ctx and truncate() for prefix reuse. - Pinned deps (cmake/deps.cmake, sha256): nlohmann/json 3.12.0, cpp-httplib 0.57.1, minja 021c2293. All MIT, credited in NOTICE. - Tests: unit tests for sampler, text stream and openai (CI); a chat-template golden against HF apply_chat_template (12 cases, CI); server_e2e.py over HTTP with Lemonade's flags. Qwen3-0.6B: greedy /v1/completions over HTTP is byte-identical to HF transformers' fp32 greedy continuation. Co-Authored-By: Claude Opus 5.5 --- CMakeLists.txt | 31 +- NOTICE | 5 + README.md | 12 +- cmake/deps.cmake | 43 +++ docs/PORTING.md | 2 +- docs/server.md | 128 +++++++++ engine/backends/cpu/cpu_model.cpp | 14 +- engine/backends/cpu/cpu_model.h | 11 +- engine/core/sampler.cpp | 75 +++++ engine/core/sampler.h | 49 ++++ engine/core/tokenizer.cpp | 4 + engine/core/tokenizer.h | 2 + engine/server/chat_template.cpp | 65 +++++ engine/server/chat_template.h | 54 ++++ engine/server/generator.cpp | 124 ++++++++ engine/server/generator.h | 78 +++++ engine/server/main.cpp | 314 +++++++++++++++++++++ engine/server/openai.cpp | 192 +++++++++++++ engine/server/openai.h | 67 +++++ engine/server/text_stream.cpp | 156 ++++++++++ engine/server/text_stream.h | 74 +++++ tests/golden/qwen3-0.6b-capitals/meta.json | 3 +- tests/golden/qwen3-template/cases.jsonl | 12 + tests/golden_template.cpp | 68 +++++ tests/server_e2e.py | 209 ++++++++++++++ tests/test_openai.cpp | 120 ++++++++ tests/test_sampler.cpp | 75 +++++ tests/test_text_stream.cpp | 122 ++++++++ tools/copyright.py | 2 +- tools/golden/make_golden.py | 2 +- tools/golden/make_template_golden.py | 91 ++++++ 31 files changed, 2185 insertions(+), 19 deletions(-) create mode 100644 cmake/deps.cmake create mode 100644 docs/server.md create mode 100644 engine/core/sampler.cpp create mode 100644 engine/core/sampler.h create mode 100644 engine/server/chat_template.cpp create mode 100644 engine/server/chat_template.h create mode 100644 engine/server/generator.cpp create mode 100644 engine/server/generator.h create mode 100644 engine/server/main.cpp create mode 100644 engine/server/openai.cpp create mode 100644 engine/server/openai.h create mode 100644 engine/server/text_stream.cpp create mode 100644 engine/server/text_stream.h create mode 100644 tests/golden/qwen3-template/cases.jsonl create mode 100644 tests/golden_template.cpp create mode 100644 tests/server_e2e.py create mode 100644 tests/test_openai.cpp create mode 100644 tests/test_sampler.cpp create mode 100644 tests/test_text_stream.cpp create mode 100644 tools/golden/make_template_golden.py diff --git a/CMakeLists.txt b/CMakeLists.txt index d779bda2..f9e02d7d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -37,7 +37,8 @@ add_library(onebit_core STATIC engine/core/model_config.cpp engine/core/parallel.cpp engine/core/unicode.cpp - engine/core/tokenizer.cpp) + engine/core/tokenizer.cpp + engine/core/sampler.cpp) target_include_directories(onebit_core PUBLIC engine/core) find_package(Threads REQUIRED) target_link_libraries(onebit_core PUBLIC Threads::Threads) @@ -46,11 +47,25 @@ add_library(onebit_cpu STATIC engine/backends/cpu/cpu_model.cpp) target_include_directories(onebit_cpu PUBLIC engine/backends/cpu) target_link_libraries(onebit_cpu PUBLIC onebit_core) +include(cmake/deps.cmake) + +add_library(onebit_server_lib STATIC + engine/server/text_stream.cpp + engine/server/chat_template.cpp + engine/server/generator.cpp + engine/server/openai.cpp) +target_include_directories(onebit_server_lib PUBLIC engine/server) +target_link_libraries(onebit_server_lib PUBLIC onebit_cpu onebit_minja) + +add_executable(onebit_server engine/server/main.cpp) +set_target_properties(onebit_server PROPERTIES OUTPUT_NAME 1bit-server) +target_link_libraries(onebit_server PRIVATE onebit_server_lib onebit_httplib) + # Unit tests: synthetic data only, so they run anywhere (including CI). -foreach(t test_gguf test_dequant test_unicode) +foreach(t test_gguf test_dequant test_unicode test_sampler test_text_stream test_openai) add_executable(${t} tests/${t}.cpp) target_include_directories(${t} PRIVATE tests) - target_link_libraries(${t} PRIVATE onebit_core) + target_link_libraries(${t} PRIVATE onebit_server_lib) add_test(NAME ${t} COMMAND ${t}) endforeach() @@ -64,6 +79,11 @@ set(ENGINE_GOLDEN_MODEL "" CACHE FILEPATH "GGUF model for the golden test") set(ENGINE_GOLDEN_DIR "" CACHE PATH "Golden logits directory for the golden test") if(ENGINE_GOLDEN_MODEL AND ENGINE_GOLDEN_DIR) add_test(NAME golden_cpu COMMAND golden_cpu ${ENGINE_GOLDEN_MODEL} ${ENGINE_GOLDEN_DIR}) + # The server end to end over HTTP, with Lemonade's launch flags. + find_package(Python3 REQUIRED COMPONENTS Interpreter) + add_test(NAME server_e2e + COMMAND ${Python3_EXECUTABLE} ${CMAKE_SOURCE_DIR}/tests/server_e2e.py --server $ + --model ${ENGINE_GOLDEN_MODEL} --golden ${ENGINE_GOLDEN_DIR}) endif() # Tokenizer golden test: encode/decode/pre-token split against HF tokenizers. @@ -71,6 +91,11 @@ endif() # (tools/golden/make_tokenizer_golden.py, docs/tokenizer.md). add_executable(golden_tokenizer tests/golden_tokenizer.cpp) target_link_libraries(golden_tokenizer PRIVATE onebit_core) +add_executable(golden_template tests/golden_template.cpp) +target_link_libraries(golden_template PRIVATE onebit_server_lib) +add_test(NAME golden_template_qwen3 + COMMAND golden_template ${CMAKE_SOURCE_DIR}/tests/golden/qwen3-tokenizer/vocab.gguf + ${CMAKE_SOURCE_DIR}/tests/golden/qwen3-template/cases.jsonl) add_test(NAME golden_tokenizer_qwen3 COMMAND golden_tokenizer ${CMAKE_SOURCE_DIR}/tests/golden/qwen3-tokenizer/vocab.gguf ${CMAKE_SOURCE_DIR}/tests/golden/qwen3-tokenizer/cases.tsv) diff --git a/NOTICE b/NOTICE index de47c8c5..ad3e1d6a 100644 --- a/NOTICE +++ b/NOTICE @@ -6,3 +6,8 @@ This product is licensed under the Apache License, Version 2.0 (see LICENSE). engine/core/unicode_tables.inc is generated from the Unicode Character Database (via Python's unicodedata module), Copyright Unicode, Inc., used under the Unicode License v3 (https://www.unicode.org/license.txt). + +Third-party components fetched at build time (cmake/deps.cmake), all MIT: + nlohmann/json 3.12.0 Copyright (c) 2013-2025 Niels Lohmann + cpp-httplib 0.57.1 Copyright (c) 2017 yhirose + google/minja (021c2293) Copyright 2024 Google LLC diff --git a/README.md b/README.md index 37d05566..d6a533ac 100644 --- a/README.md +++ b/README.md @@ -21,10 +21,12 @@ XDNA 2 NPU, Radeon iGPU (Vulkan and HRX), and a CPU reference path. It serves th llama-server-compatible OpenAI HTTP API, so [Lemonade](https://github.com/lemonade-sdk/lemonade) and any OpenAI client can drive it as a drop-in backend. -> **Status: scaffold.** Nothing runs yet. Components land one at a time, each with a -> test against the CPU reference (see [CONTRIBUTING.md](CONTRIBUTING.md)). The -> development history lives in [1bit-MONSTER](https://github.com/1bit-MONSTER/1bit-MONSTER); -> this repository contains only code that has been verified. +> **Status:** the CPU reference, tokenizer and server work end to end on Qwen3 +> (`1bit-server`, see [docs/server.md](docs/server.md)). The NPU and GPU backends +> come next. Components land one at a time, each with a test against the CPU +> reference (see [CONTRIBUTING.md](CONTRIBUTING.md)). The development history +> lives in [1bit-MONSTER](https://github.com/1bit-MONSTER/1bit-MONSTER); this +> repository contains only code that has been verified. ## Layout @@ -35,7 +37,7 @@ and any OpenAI client can drive it as a drop-in backend. | `engine/backends/npu/` | XDNA 2 via XRT, full-ELF kernels only (no xclbins) | | `engine/backends/gpu/` | llama.cpp (AMD-Ecosystem fork) with the Vulkan and HRX20 devices | | `engine/backends/cpu/` | fp32 reference forward pass; the correctness oracle | -| `engine/server/` | llama-server-compatible HTTP: `/v1/chat/completions`, `/v1/completions`, `/v1/models`, `/health` | +| `engine/server/` | `1bit-server`: Lemonade's backend protocol, i.e. `/health`, `/v1/models`, `/v1/chat/completions`, `/v1/completions` | | `tests/` | Golden-logit tests per backend against the CPU reference | | `lemonade/` | Upstream patch: `BackendDescriptor` + `WrappedServer` for this engine | diff --git a/cmake/deps.cmake b/cmake/deps.cmake new file mode 100644 index 00000000..50c61b8c --- /dev/null +++ b/cmake/deps.cmake @@ -0,0 +1,43 @@ +# Copyright 2026 bong-water-water-bong +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Third-party dependencies, each pinned to a release (or commit) and its +# sha256. All are header-only and MIT-licensed; see NOTICE. +include(FetchContent) + +FetchContent_Declare(nlohmann_json + URL https://github.com/nlohmann/json/releases/download/v3.12.0/json.tar.xz + URL_HASH SHA256=42f6e95cad6ec532fd372391373363b62a14af6d771056dbfc86160e6dfff7aa) +set(JSON_BuildTests OFF CACHE INTERNAL "") +set(JSON_Install OFF CACHE INTERNAL "") + +# SOURCE_SUBDIR points at a directory that does not exist, so FetchContent +# only downloads these and does not run their own CMake projects. +FetchContent_Declare(cpp_httplib + URL https://github.com/yhirose/cpp-httplib/archive/refs/tags/v0.57.1.tar.gz + URL_HASH SHA256=5c9e56b6638eb415dc451ac57531e84931fcb50e261ca3dd74eec6b484fb4712 + SOURCE_SUBDIR _none) +FetchContent_Declare(minja + URL https://github.com/google/minja/archive/021c2293c187789ef13d56c6cfd89c9b134fd80f.tar.gz + URL_HASH SHA256=dc3ddd37497b79a4cd35fd41550e22b3b0139439de8be6eab3d2d024fed47bb4 + SOURCE_SUBDIR _none) + +FetchContent_MakeAvailable(nlohmann_json cpp_httplib minja) + +add_library(onebit_httplib INTERFACE) +target_include_directories(onebit_httplib SYSTEM INTERFACE ${cpp_httplib_SOURCE_DIR}) +add_library(onebit_minja INTERFACE) +target_include_directories(onebit_minja SYSTEM INTERFACE ${minja_SOURCE_DIR}/include) +target_link_libraries(onebit_minja INTERFACE nlohmann_json::nlohmann_json) diff --git a/docs/PORTING.md b/docs/PORTING.md index 772b5669..f3429232 100644 --- a/docs/PORTING.md +++ b/docs/PORTING.md @@ -28,7 +28,7 @@ and what has to be true before it lands here. Refs are branches or commits in th | NPU 16-tile layer kernel | branch `bench/fastlane-16tile-corrections-2026-09-22` (435bf36e7) | 24/24 tokens, summed KL 0.000466, 0.340 ms/layer | rebuild from source, reproduce | | GPU backend | AMD-Ecosystem/llama.cpp fork, `GGML_VULKAN` + `GGML_HRX2`; recipe on `fix/zaya-lmhead-evidence` | Vulkan0 74.8 tok/s, HRX20 18.4 on zaya1-8b; Q4NX is HRX20-only | pinned submodule, linked (no dlopen of copied structs) | | Router | `src/model_router.cpp` (Q4NX rule) + Laya scorer, branch `backup/laya-and-results-2026-09-22` | Laya not yet checked against its Python reference | Laya matches its Python reference within a stated tolerance | -| Server | `src/server/` | works | llama-server flag and endpoint compatibility test | +| Server | `src/server/` | **landed** as a new implementation (docs/server.md) | Lemonade's flags and endpoints, end-to-end test | | Lemonade recipe | lemonade `src/cpp/include/lemon/backends/*` pattern | n/a | loads Qwen3-0.6B through a local Lemonade build | Known traps carried over: diff --git a/docs/server.md b/docs/server.md new file mode 100644 index 00000000..6b03933b --- /dev/null +++ b/docs/server.md @@ -0,0 +1,128 @@ + +# Server + +`1bit-server` is the engine's HTTP front end. It speaks the backend protocol +Lemonade uses for its GGUF recipes (`llamacpp`, `llamacpp-hrx`): the launch +flags, a `/health` probe, and the OpenAI endpoints with llama-server's +`timings` extension. Lemonade can run it wherever it runs those backends today. + +Only the protocol is shared with llama-server. The inference is this repository's own: + +- model code (`engine/backends`) +- tokenizer, sampler and chat-template rendering (`engine/core`, `engine/server`) + +## Launch + +```bash +1bit-server -m model.gguf --ctx-size 4096 --device cpu --port 8080 --jinja --metrics --parallel 1 +``` + +| Flag | Meaning | +|---|---| +| `-m, --model` | GGUF model (required) | +| `-c, --ctx-size` | context length (default: min(training context, 4096)) | +| `--host`, `--port` | listen address (default `127.0.0.1:8080`) | +| `--device` | `auto` or `cpu`. More devices are added with each backend; an unknown device is an error, never a silent fallback | +| `-t, --threads` | CPU threads | +| `-a, --alias` | model id reported by the API | +| `-np, --parallel` | only `1` | +| `--jinja`, `--metrics` | accepted for compatibility. The GGUF chat template is always used; there is no `/metrics` yet | + +Unknown flags are rejected, so a misconfigured launch fails loudly. + +## Endpoints + +| Endpoint | Notes | +|---|---| +| `GET /health` | `503 {"error":{"message":"Loading model"}}` while loading, then `200 {"status":"ok"}` | +| `GET /v1/models` | one model, id = alias | +| `POST /v1/chat/completions` | messages rendered with the GGUF's Jinja template; streaming or not | +| `POST /v1/completions` | raw prompt; streaming or not | + +Every response carries OpenAI `usage` (including +`prompt_tokens_details.cached_tokens`) and llama-server `timings` (`prompt_n`, +`prompt_ms`, `predicted_n`, `predicted_per_second`, `cache_n` and others). These +are the fields Lemonade's telemetry reads. Streams end with `data: [DONE]`. The +final chunk carries `finish_reason` and `timings`, and +`stream_options.include_usage` adds a usage chunk. + +### Request fields + +- **Sampling:** `temperature` (0 = greedy), `top_p`, `top_k`, `min_p` and `seed`. + The defaults match llama-server: 0.8, 0.95, 40 and 0.05, with a random seed. +- **Limits and streaming:** `max_completion_tokens` / `max_tokens` / `n_predict`, + `stop` (a string or an array), `stream` and `stream_options`. +- **Chat:** `chat_template_kwargs` (for example `{"enable_thinking": false}`), + `reasoning_format` (`none` keeps `` in `content`), and `tools` + (rendered into the prompt). +- **Other fields** are ignored. +- **Rejected with 400**, because they would otherwise be silently wrong: + `n > 1`, non-text content parts, batched prompts, and a prompt that does not + fit the context (`exceed_context_size_error`). + +### Reasoning + +A leading `...` block goes to `message.reasoning_content` (in +streams, `delta.reasoning_content`), and the answer goes to `content`. This is +the split Lemonade's chat, MCP and Ollama paths read. When the template itself +opens the block (the prompt ends in ``), the output starts in reasoning. + +### Prompt cache + +The previous request's tokens stay in the KV cache. A new request reuses the +longest shared prefix and evaluates only the rest, reported as `cache_n`. That +makes multi-turn chat cost one turn, not the whole conversation. + +## Verification + +`tests/server_e2e.py` launches the binary with Lemonade's flags and checks the +contract over HTTP. It runs in `ctest` when a golden model is configured. On +Qwen3-0.6B it checks: + +- **Correctness:** the greedy `/v1/completions` text is byte-identical to HF + transformers' fp32 greedy continuation from the golden, and `prompt_tokens` + matches HF's tokenization. +- **Streaming:** + - streamed text equals non-streamed text, for completions and for chat + (both `content` and `reasoning_content`) + - there is exactly one final chunk, with `finish_reason` and `timings` + - the usage chunk appears + - the stream ends with `[DONE]` +- **Prefix cache:** a repeated request reuses the whole prompt but its last + token (`cache_n = n_prompt - 1`) and gives the same text. +- **Stops and chat:** + - stop strings truncate before the match, with `finish_reason: stop` + - chat stops on the end-of-turn token + - `enable_thinking: false` gives plain content + - thinking on splits the reasoning out +- **Errors:** bad JSON, `n = 2` and an over-long prompt give 400; an unknown route gives 404. + +Unit tests (CI): `test_text_stream` covers UTF-8 split across tokens, stop +strings across tokens and reasoning routing byte by byte. `test_openai` covers +request parsing and response shapes, checked structurally. `test_sampler` +covers greedy, seeding, top-k, top-p and min-p. `golden_template_qwen3` +renders 12 conversations byte-identical to HF `apply_chat_template`, including +tools, tool calls and the thinking toggle. + +## Not yet + +- **Tool-call output parsing.** Tools are rendered into the prompt, but a + generated `` block is returned as text, not as `tool_calls`. +- **Endpoints:** `/v1/responses`, `/v1/embeddings`, `/tokenize`, `/slots` and `/metrics`. +- **Concurrency:** more than one concurrent sequence. Requests queue. +- **Devices:** anything but the CPU reference. The NPU and GPU backends plug in behind the same server. diff --git a/engine/backends/cpu/cpu_model.cpp b/engine/backends/cpu/cpu_model.cpp index 2ab05e12..ce6d170d 100644 --- a/engine/backends/cpu/cpu_model.cpp +++ b/engine/backends/cpu/cpu_model.cpp @@ -26,9 +26,8 @@ namespace onebit { namespace { -// KV cache length. Kept well under the training context so the reference -// stays small; raise it when a test needs longer sequences. -constexpr uint32_t kMaxCtx = 4096; +// Default KV cache length when the caller does not choose one. +constexpr uint32_t kDefaultCtx = 4096; void rms_norm(const float* x, const float* w, float* out, uint32_t n, float eps) { double ss = 0.0; @@ -87,7 +86,8 @@ class Loader { } // namespace -std::expected CpuModel::load(const std::string& gguf_path, size_t n_threads) { +std::expected CpuModel::load(const std::string& gguf_path, size_t n_threads, + uint32_t n_ctx) { auto file = GgufFile::open(gguf_path); if (!file) return std::unexpected(file.error()); auto cfg = read_model_config(*file); @@ -130,7 +130,9 @@ std::expected CpuModel::load(const std::string& gguf_path #undef TAKE if (auto r = ld.check_all_used(); !r) return std::unexpected(r.error()); - m.n_ctx_ = std::min(c.n_ctx_train, kMaxCtx); + if (n_ctx > c.n_ctx_train) + return std::unexpected(std::format("n_ctx {} exceeds the model's training context {}", n_ctx, c.n_ctx_train)); + m.n_ctx_ = n_ctx ? n_ctx : std::min(c.n_ctx_train, kDefaultCtx); const size_t cache = size_t{c.n_layer} * m.n_ctx_ * KV; m.k_cache_.assign(cache, 0.0f); m.v_cache_.assign(cache, 0.0f); @@ -150,6 +152,8 @@ std::expected CpuModel::load(const std::string& gguf_path void CpuModel::reset() { n_past_ = 0; } +void CpuModel::truncate(uint32_t n) { n_past_ = std::min(n, n_past_); } + void CpuModel::matvec(const std::vector& w, const float* x, float* y, uint32_t rows, uint32_t cols) { pool_->parallel_for(rows, [&](size_t begin, size_t end) { for (size_t r = begin; r < end; ++r) { diff --git a/engine/backends/cpu/cpu_model.h b/engine/backends/cpu/cpu_model.h index c86191d4..7beb4135 100644 --- a/engine/backends/cpu/cpu_model.h +++ b/engine/backends/cpu/cpu_model.h @@ -36,15 +36,22 @@ class CpuModel { // Loads every tensor of the GGUF. Fails if a tensor is missing, has the // wrong shape, has an unsupported type, or is present but unused (an // unused tensor means the model has semantics this code does not model). - static std::expected load(const std::string& gguf_path, size_t n_threads = 0); + // n_ctx: KV cache length; 0 uses min(training context, 4096). + static std::expected load(const std::string& gguf_path, size_t n_threads = 0, + uint32_t n_ctx = 0); const ModelConfig& config() const { return cfg_; } // Clears the KV cache. void reset(); - // Number of tokens in the KV cache. + // Number of tokens in the KV cache, and its capacity. uint32_t n_past() const { return n_past_; } + uint32_t n_ctx() const { return n_ctx_; } + + // Drops cached positions from n onwards (n <= n_past()), so a new + // sequence sharing the first n tokens continues from there. + void truncate(uint32_t n); // Runs one token at position n_past() and returns the next-token logits // (config().n_vocab floats). Fails when the context is full. diff --git a/engine/core/sampler.cpp b/engine/core/sampler.cpp new file mode 100644 index 00000000..70854c18 --- /dev/null +++ b/engine/core/sampler.cpp @@ -0,0 +1,75 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "sampler.h" + +#include +#include +#include +#include + +namespace onebit { + +int32_t argmax(std::span logits) { + return int32_t(std::max_element(logits.begin(), logits.end()) - logits.begin()); +} + +int32_t Sampler::sample(std::span logits) { + if (p_.temperature <= 0.0f || p_.top_k == 1) return argmax(logits); + + // Candidates sorted by logit, descending (ties: lower id first). + std::vector ids(logits.size()); + std::iota(ids.begin(), ids.end(), 0); + size_t k = ids.size(); + if (p_.top_k > 0 && size_t(p_.top_k) < k) k = size_t(p_.top_k); + auto by_logit = [&](int32_t a, int32_t b) { return logits[a] > logits[b] || (logits[a] == logits[b] && a < b); }; + std::partial_sort(ids.begin(), ids.begin() + std::ptrdiff_t(k), ids.end(), by_logit); + ids.resize(k); + + // Untempered probabilities of the candidates. + const double mx = logits[ids[0]]; + std::vector prob(k); + double sum = 0.0; + for (size_t i = 0; i < k; ++i) sum += prob[i] = std::exp(double(logits[ids[i]]) - mx); + for (double& q : prob) q /= sum; + + // top-p: keep the smallest prefix whose mass reaches top_p. + if (p_.top_p < 1.0f) { + double cum = 0.0; + size_t keep = k; + for (size_t i = 0; i < k; ++i) { + cum += prob[i]; + if (cum >= p_.top_p) { + keep = i + 1; + break; + } + } + k = keep; + } + // min-p: drop candidates below min_p times the top probability. + if (p_.min_p > 0.0f) { + const double floor = p_.min_p * prob[0]; + size_t keep = 1; + while (keep < k && prob[keep] >= floor) ++keep; + k = keep; + } + + // Temperature, then draw. + std::vector w(k); + for (size_t i = 0; i < k; ++i) w[i] = std::exp((double(logits[ids[i]]) - mx) / p_.temperature); + std::discrete_distribution dist(w.begin(), w.end()); + return ids[dist(rng_)]; +} + +} // namespace onebit diff --git a/engine/core/sampler.h b/engine/core/sampler.h new file mode 100644 index 00000000..251f3825 --- /dev/null +++ b/engine/core/sampler.h @@ -0,0 +1,49 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// sampler.h: next-token selection from logits. +// +// Order follows llama.cpp's default chain: top-k, top-p and min-p filter the +// candidates on the untempered distribution, then temperature is applied and +// one token is drawn. temperature <= 0 selects the argmax (greedy). +#pragma once + +#include +#include +#include + +namespace onebit { + +struct SamplingParams { + float temperature = 0.8f; + int32_t top_k = 40; // <= 0 disables + float top_p = 0.95f; // >= 1 disables + float min_p = 0.05f; // <= 0 disables + uint64_t seed = 0; +}; + +class Sampler { +public: + explicit Sampler(const SamplingParams& p) : p_(p), rng_(p.seed) {} + + int32_t sample(std::span logits); + +private: + SamplingParams p_; + std::mt19937_64 rng_; +}; + +int32_t argmax(std::span logits); + +} // namespace onebit diff --git a/engine/core/tokenizer.cpp b/engine/core/tokenizer.cpp index a611cc01..e64e6bbb 100644 --- a/engine/core/tokenizer.cpp +++ b/engine/core/tokenizer.cpp @@ -254,6 +254,10 @@ std::expected, std::string> Tokenizer::encode(std::string_v return out; } +bool Tokenizer::is_control(int32_t id) const { + return id >= 0 && size_t(id) < token_type_.size() && token_type_[size_t(id)] == kTypeControl; +} + std::expected Tokenizer::decode(std::span ids, bool skip_special) const { std::string out; for (int32_t id : ids) { diff --git a/engine/core/tokenizer.h b/engine/core/tokenizer.h index 74dc516b..f82140a7 100644 --- a/engine/core/tokenizer.h +++ b/engine/core/tokenizer.h @@ -50,6 +50,8 @@ class Tokenizer { int32_t bos_id() const { return bos_; } int32_t eos_id() const { return eos_; } const std::string& token_text(int32_t id) const { return tokens_.at(size_t(id)); } + // CONTROL tokens (<|im_end|>, <|endoftext|>, ...): never rendered as text by default. + bool is_control(int32_t id) const; private: Tokenizer() = default; diff --git a/engine/server/chat_template.cpp b/engine/server/chat_template.cpp new file mode 100644 index 00000000..a79d6eff --- /dev/null +++ b/engine/server/chat_template.cpp @@ -0,0 +1,65 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "chat_template.h" + +#include + +namespace onebit { + +ChatTemplate::ChatTemplate(ChatTemplate&&) noexcept = default; +ChatTemplate& ChatTemplate::operator=(ChatTemplate&&) noexcept = default; +ChatTemplate::~ChatTemplate() = default; + +std::expected ChatTemplate::from_source(const std::string& source, const std::string& bos, + const std::string& eos) { + ChatTemplate t; + try { + t.impl_ = std::make_unique(source, bos, eos); + } catch (const std::exception& e) { + return std::unexpected(std::string("chat template does not parse: ") + e.what()); + } + return t; +} + +std::expected ChatTemplate::from_gguf(const GgufFile& f) { + auto source = f.get_string("tokenizer.chat_template"); + if (!source) return std::unexpected(std::string("GGUF has no tokenizer.chat_template")); + auto token_text = [&](const char* key) -> std::string { + auto id = f.get_uint(key); + const GgufArray* tokens = f.get_array("tokenizer.ggml.tokens"); + if (!id || !tokens || *id >= tokens->size()) return ""; + const auto* s = std::get_if(&(*tokens)[*id].v); + return s ? *s : ""; + }; + return from_source(*source, token_text("tokenizer.ggml.bos_token_id"), token_text("tokenizer.ggml.eos_token_id")); +} + +std::expected ChatTemplate::render(const nlohmann::ordered_json& messages, + const nlohmann::ordered_json& tools, + bool add_generation_prompt, + const nlohmann::ordered_json& kwargs) const { + minja::chat_template_inputs in; + in.messages = messages; + in.tools = tools; + in.add_generation_prompt = add_generation_prompt; + in.extra_context = kwargs.is_object() ? kwargs : nlohmann::ordered_json::object(); + try { + return impl_->apply(in); + } catch (const std::exception& e) { + return std::unexpected(std::string("chat template failed: ") + e.what()); + } +} + +} // namespace onebit diff --git a/engine/server/chat_template.h b/engine/server/chat_template.h new file mode 100644 index 00000000..1aa887a9 --- /dev/null +++ b/engine/server/chat_template.h @@ -0,0 +1,54 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// chat_template.h: renders OpenAI-style messages with the model's own Jinja +// chat template (GGUF key tokenizer.chat_template), via minja. +#pragma once + +#include +#include +#include + +#include + +#include "gguf.h" + +namespace minja { +class chat_template; +} + +namespace onebit { + +class ChatTemplate { +public: + static std::expected from_gguf(const GgufFile& f); + static std::expected from_source(const std::string& source, const std::string& bos, + const std::string& eos); + + // messages: array of {role, content, ...}; tools: array or null. + // kwargs: extra template variables (chat_template_kwargs, e.g. enable_thinking). + std::expected render(const nlohmann::ordered_json& messages, + const nlohmann::ordered_json& tools, bool add_generation_prompt, + const nlohmann::ordered_json& kwargs) const; + + ChatTemplate(ChatTemplate&&) noexcept; + ChatTemplate& operator=(ChatTemplate&&) noexcept; + ~ChatTemplate(); + +private: + ChatTemplate() = default; + std::unique_ptr impl_; +}; + +} // namespace onebit diff --git a/engine/server/generator.cpp b/engine/server/generator.cpp new file mode 100644 index 00000000..ba7eb2bb --- /dev/null +++ b/engine/server/generator.cpp @@ -0,0 +1,124 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "generator.h" + +#include +#include +#include + +namespace onebit { + +namespace { + +// Control tokens that end a turn in the chat formats of supported models. +// tokenizer.ggml.eos_token_id is always included as well. +constexpr const char* kEndOfTurn[] = {"<|im_end|>", "<|endoftext|>", "<|eot_id|>", "<|end_of_text|>", + "<|end|>", ""}; + +double ms_since(std::chrono::steady_clock::time_point t) { + return std::chrono::duration(std::chrono::steady_clock::now() - t).count(); +} + +} // namespace + +Generator::Generator(CpuModel model, Tokenizer tokenizer) : model_(std::move(model)), tok_(std::move(tokenizer)) { + if (tok_.eos_id() >= 0) eog_.insert(tok_.eos_id()); + for (int32_t id = 0; size_t(id) < tok_.n_vocab(); ++id) { + if (!tok_.is_control(id)) continue; + for (const char* t : kEndOfTurn) + if (tok_.token_text(id) == t) eog_.insert(id); + } +} + +std::expected Generator::generate(const GenerateRequest& req, const OnDelta& on_delta) { + std::lock_guard lock(mu_); + if (req.prompt.empty()) return std::unexpected(std::string("prompt is empty")); + if (req.prompt.size() >= model_.n_ctx()) + return std::unexpected(std::format("prompt is {} tokens; the context holds {}", req.prompt.size(), model_.n_ctx())); + + GenerateResult r; + r.prompt_tokens = int32_t(req.prompt.size()); + + // Reuse the shared prefix; the last prompt token is always re-evaluated + // because its logits are needed to sample the first new token. + size_t keep = 0; + while (keep < cached_.size() && keep < req.prompt.size() && cached_[keep] == req.prompt[keep]) ++keep; + keep = std::min(keep, req.prompt.size() - 1); + model_.truncate(uint32_t(keep)); + cached_.resize(keep); + r.cache_n = int32_t(keep); + + auto t0 = std::chrono::steady_clock::now(); + const std::vector* logits = nullptr; + for (size_t i = keep; i < req.prompt.size(); ++i) { + auto out = model_.forward(req.prompt[i]); + if (!out) { + cached_.clear(); + model_.reset(); + return std::unexpected(out.error()); + } + cached_.push_back(req.prompt[i]); + logits = *out; + } + r.prompt_ms = ms_since(t0); + + Sampler sampler(req.sampling); + TextStream stream(req.stops, req.split_reasoning, req.reasoning_open); + auto t1 = std::chrono::steady_clock::now(); + bool cancelled = false; + r.finish_reason = "length"; + for (;;) { + if (req.max_tokens >= 0 && r.predicted_n >= req.max_tokens) break; + const int32_t id = sampler.sample(*logits); + if (is_end_of_generation(id)) { + r.finish_reason = "stop"; + break; + } + ++r.predicted_n; + const int32_t one[1] = {id}; + auto piece = tok_.decode(one, /*skip_special=*/true); + if (!piece) return std::unexpected(piece.error()); + TextDelta d = stream.push(*piece); + if (!d.empty() && on_delta && !on_delta(d)) { + cancelled = true; + break; + } + if (stream.stopped()) { + r.finish_reason = "stop"; + break; + } + if (model_.n_past() >= model_.n_ctx()) break; // context full + auto out = model_.forward(id); + if (!out) { + cached_.clear(); + model_.reset(); + return std::unexpected(out.error()); + } + cached_.push_back(id); + logits = *out; + } + if (!cancelled) { + TextDelta d = stream.finish(); + if (!d.empty() && on_delta) on_delta(d); + } else { + r.finish_reason = "stop"; + } + r.predicted_ms = ms_since(t1); + r.content = stream.content(); + r.reasoning = stream.reasoning(); + return r; +} + +} // namespace onebit diff --git a/engine/server/generator.h b/engine/server/generator.h new file mode 100644 index 00000000..d5c186d4 --- /dev/null +++ b/engine/server/generator.h @@ -0,0 +1,78 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// generator.h: runs one generation at a time over a model and tokenizer. +// +// Keeps the previous request's tokens in the KV cache and reuses the longest +// shared prefix (llama-server's cache_n), which is what makes multi-turn chat +// cheap: only the new turn is evaluated. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "cpu_model.h" +#include "sampler.h" +#include "text_stream.h" +#include "tokenizer.h" + +namespace onebit { + +struct GenerateRequest { + std::vector prompt; + SamplingParams sampling; + int32_t max_tokens = -1; // < 0: until end of generation or a full context + std::vector stops; + bool split_reasoning = false; + bool reasoning_open = false; +}; + +struct GenerateResult { + std::string content, reasoning; + std::string finish_reason; // "stop" or "length" + int32_t prompt_tokens = 0; // whole prompt, including the reused prefix + int32_t cache_n = 0; // prompt tokens reused from the previous request + int32_t predicted_n = 0; + double prompt_ms = 0.0, predicted_ms = 0.0; +}; + +// Called with each delta as it becomes sendable. Return false to cancel +// (client went away); the result then has finish_reason "stop". +using OnDelta = std::function; + +class Generator { +public: + Generator(CpuModel model, Tokenizer tokenizer); + + const Tokenizer& tokenizer() const { return tok_; } + uint32_t n_ctx() const { return model_.n_ctx(); } + + std::expected generate(const GenerateRequest& req, const OnDelta& on_delta); + +private: + bool is_end_of_generation(int32_t id) const { return eog_.contains(id); } + + std::mutex mu_; + CpuModel model_; + Tokenizer tok_; + std::set eog_; + std::vector cached_; // tokens currently in the KV cache +}; + +} // namespace onebit diff --git a/engine/server/main.cpp b/engine/server/main.cpp new file mode 100644 index 00000000..29437b06 --- /dev/null +++ b/engine/server/main.cpp @@ -0,0 +1,314 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// 1bit-server: llama-server-compatible OpenAI HTTP server. +// +// Speaks the subset of llama-server's interface that Lemonade drives: the +// launch flags it passes, GET /health (503 while loading), /v1/models, +// /v1/chat/completions and /v1/completions (streaming and not), with both +// OpenAI `usage` and llama-server `timings` in the responses. +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "chat_template.h" +#include "generator.h" +#include "openai.h" + +using namespace onebit; +using json = openai::json; + +namespace { + +constexpr const char* kVersion = "0.0.1"; + +struct Args { + std::string model; + std::string host = "127.0.0.1"; + int port = 8080; + uint32_t ctx_size = 0; + size_t threads = 0; + std::string device = "auto"; + std::string alias; +}; + +void usage(FILE* out) { + std::fprintf(out, + "usage: 1bit-server -m MODEL.gguf [options]\n" + " -m, --model PATH GGUF model (required)\n" + " -c, --ctx-size N context length (default: min(training context, 4096))\n" + " --host HOST listen address (default 127.0.0.1)\n" + " --port N listen port (default 8080)\n" + " --device NAME auto | cpu (default auto)\n" + " -t, --threads N CPU threads (default: all)\n" + " -a, --alias NAME model name reported by the API (default: file name)\n" + " -np, --parallel N concurrent sequences; only 1 is supported\n" + " --jinja accepted: the GGUF chat template is always used\n" + " --metrics accepted: no /metrics endpoint yet\n" + " -h, --help, --version\n"); +} + +std::optional parse_args(int argc, char** argv) { + Args a; + auto need = [&](int& i) -> const char* { + if (i + 1 >= argc) { + std::fprintf(stderr, "error: %s needs a value\n", argv[i]); + return nullptr; + } + return argv[++i]; + }; + auto to_int = [](const char* s, long long lo, long long hi, long long& out) { + char* end = nullptr; + out = std::strtoll(s, &end, 10); + return end && *end == '\0' && out >= lo && out <= hi; + }; + for (int i = 1; i < argc; ++i) { + const std::string f = argv[i]; + long long n = 0; + const char* v = nullptr; + if (f == "-h" || f == "--help") { + usage(stdout); + std::exit(0); + } else if (f == "--version") { + std::printf("1bit-server %s\n", kVersion); + std::exit(0); + } else if (f == "-m" || f == "--model") { + if (!(v = need(i))) return std::nullopt; + a.model = v; + } else if (f == "-c" || f == "--ctx-size") { + if (!(v = need(i)) || !to_int(v, 0, UINT32_MAX, n)) return std::fprintf(stderr, "error: bad --ctx-size\n"), std::nullopt; + a.ctx_size = uint32_t(n); + } else if (f == "--host") { + if (!(v = need(i))) return std::nullopt; + a.host = v; + } else if (f == "--port") { + if (!(v = need(i)) || !to_int(v, 0, 65535, n)) return std::fprintf(stderr, "error: bad --port\n"), std::nullopt; + a.port = int(n); + } else if (f == "--device" || f == "-dev") { + if (!(v = need(i))) return std::nullopt; + a.device = v; + } else if (f == "-t" || f == "--threads") { + if (!(v = need(i)) || !to_int(v, 0, 4096, n)) return std::fprintf(stderr, "error: bad --threads\n"), std::nullopt; + a.threads = size_t(n); + } else if (f == "-a" || f == "--alias") { + if (!(v = need(i))) return std::nullopt; + a.alias = v; + } else if (f == "-np" || f == "--parallel") { + if (!(v = need(i)) || !to_int(v, 1, 1, n)) return std::fprintf(stderr, "error: only --parallel 1 is supported\n"), std::nullopt; + } else if (f == "--jinja" || f == "--metrics") { + // Accepted for llama-server compatibility; see usage(). + } else { + std::fprintf(stderr, "error: unknown argument %s\n", f.c_str()); + usage(stderr); + return std::nullopt; + } + } + if (a.model.empty()) { + std::fprintf(stderr, "error: --model is required\n"); + return std::nullopt; + } + std::string dev = a.device; + for (char& c : dev) c = char(std::tolower(static_cast(c))); + if (dev != "auto" && dev != "cpu") { + std::fprintf(stderr, "error: device '%s' is not available in this build (available: auto, cpu)\n", a.device.c_str()); + return std::nullopt; + } + if (a.alias.empty()) a.alias = std::filesystem::path(a.model).stem().string(); + return a; +} + +std::string random_id(const char* prefix) { + static std::mutex mu; + static std::mt19937_64 rng{std::random_device{}()}; + std::lock_guard lock(mu); + return std::format("{}{:016x}", prefix, rng()); +} + +bool sse(httplib::DataSink& sink, const json& j) { + const std::string s = "data: " + j.dump() + "\n\n"; + return sink.is_writable() && sink.write(s.data(), s.size()); +} + +// Everything that exists once the model has loaded. +struct Engine { + std::unique_ptr gen; + std::unique_ptr tmpl; +}; + +} // namespace + +int main(int argc, char** argv) { + auto args = parse_args(argc, argv); + if (!args) return 2; + + enum { kLoading, kReady, kFailed }; + std::atomic state{kLoading}; + std::string load_error; + Engine engine; + + httplib::Server svr; + auto send_json = [](httplib::Response& res, int status, const json& j) { + res.status = status; + res.set_content(j.dump(), "application/json"); + }; + auto ready_or_503 = [&](httplib::Response& res) { + const int s = state.load(); + if (s == kReady) return true; + if (s == kLoading) send_json(res, 503, openai::error(503, "Loading model", "unavailable_error")); + else send_json(res, 500, openai::error(500, "model failed to load: " + load_error, "server_error")); + return false; + }; + + svr.Get("/health", [&](const httplib::Request&, httplib::Response& res) { + if (ready_or_503(res)) send_json(res, 200, {{"status", "ok"}}); + }); + + auto models = [&](const httplib::Request&, httplib::Response& res) { + json m = {{"id", args->alias}, {"object", "model"}, {"created", int64_t(std::time(nullptr))}, {"owned_by", "1bit"}}; + send_json(res, 200, {{"object", "list"}, {"data", json::array({m})}}); + }; + svr.Get("/v1/models", models); + svr.Get("/models", models); + + // Shared by both endpoints: parse, build the prompt, run, respond. + auto handle = [&](const httplib::Request& req, httplib::Response& res, bool chat) { + if (!ready_or_503(res)) return; + json body = json::parse(req.body, nullptr, /*allow_exceptions=*/false); + if (body.is_discarded()) return send_json(res, 400, openai::error(400, "request body is not valid JSON", "invalid_request_error")); + auto parsed = chat ? openai::parse_chat(body) : openai::parse_completion(body); + if (!parsed) return send_json(res, 400, openai::error(400, parsed.error(), "invalid_request_error")); + openai::Request r = std::move(*parsed); + + std::string prompt = r.prompt; + if (chat) { + auto rendered = engine.tmpl->render(r.messages, r.tools, /*add_generation_prompt=*/true, r.template_kwargs); + if (!rendered) return send_json(res, 400, openai::error(400, rendered.error(), "invalid_request_error")); + prompt = std::move(*rendered); + r.gen.split_reasoning = r.split_reasoning; + r.gen.reasoning_open = prompt.ends_with("\n") || prompt.ends_with(""); + } + auto tokens = engine.gen->tokenizer().encode(prompt, /*parse_special=*/true); + if (!tokens) return send_json(res, 400, openai::error(400, tokens.error(), "invalid_request_error")); + if (tokens->empty()) return send_json(res, 400, openai::error(400, "prompt is empty", "invalid_request_error")); + if (tokens->size() >= engine.gen->n_ctx()) + return send_json(res, 400, openai::error(400, std::format("the prompt is {} tokens, which exceeds the context size of {}", + tokens->size(), engine.gen->n_ctx()), + "exceed_context_size_error")); + r.gen.prompt = std::move(*tokens); + const openai::Meta meta{random_id(chat ? "chatcmpl-" : "cmpl-"), args->alias, int64_t(std::time(nullptr))}; + + auto log = [](const GenerateResult& g) { + std::fprintf(stderr, "request: prompt %d (cached %d) in %.0f ms, generated %d in %.0f ms (%.2f tok/s), %s\n", + g.prompt_tokens, g.cache_n, g.prompt_ms, g.predicted_n, g.predicted_ms, + g.predicted_ms > 0 ? 1000.0 * g.predicted_n / g.predicted_ms : 0.0, g.finish_reason.c_str()); + }; + + if (!r.stream) { + auto g = engine.gen->generate(r.gen, {}); + if (!g) return send_json(res, 500, openai::error(500, g.error(), "server_error")); + log(*g); + return send_json(res, 200, chat ? openai::chat_response(meta, *g) : openai::completion_response(meta, *g)); + } + + res.set_header("Cache-Control", "no-cache"); + res.set_chunked_content_provider( + "text/event-stream", [&, r = std::move(r), meta, chat, log](size_t, httplib::DataSink& sink) { + if (chat && !sse(sink, openai::chat_chunk(meta, {{"role", "assistant"}, {"content", nullptr}}, nullptr))) + return false; + auto g = engine.gen->generate(r.gen, [&](const TextDelta& d) { + if (!chat) return sse(sink, openai::completion_chunk(meta, d.content, nullptr)); + json delta = json::object(); + if (!d.reasoning.empty()) delta["reasoning_content"] = d.reasoning; + if (!d.content.empty()) delta["content"] = d.content; + return sse(sink, openai::chat_chunk(meta, delta, nullptr)); + }); + if (!g) { + // In-stream error: an error object and no [DONE], as OpenAI does. + sse(sink, openai::error(500, g.error(), "server_error")); + sink.done(); + return true; + } + log(*g); + json last = chat ? openai::chat_chunk(meta, json::object(), g->finish_reason) + : openai::completion_chunk(meta, "", g->finish_reason); + last["timings"] = openai::timings(*g); + sse(sink, last); + if (r.include_usage) { + json u = chat ? openai::chat_chunk(meta, json::object(), nullptr) + : openai::completion_chunk(meta, "", nullptr); + u["choices"] = json::array(); + u["usage"] = openai::usage(*g); + sse(sink, u); + } + const std::string done = "data: [DONE]\n\n"; + if (sink.is_writable()) sink.write(done.data(), done.size()); + sink.done(); + return true; + }); + }; + for (const char* path : {"/v1/chat/completions", "/chat/completions"}) + svr.Post(path, [&](const httplib::Request& q, httplib::Response& s) { handle(q, s, true); }); + for (const char* path : {"/v1/completions", "/completions"}) + svr.Post(path, [&](const httplib::Request& q, httplib::Response& s) { handle(q, s, false); }); + + svr.set_error_handler([&](const httplib::Request&, httplib::Response& res) { + if (res.body.empty()) send_json(res, res.status, openai::error(res.status, "not found", "not_found_error")); + }); + + if (!svr.bind_to_port(args->host, args->port)) { + std::fprintf(stderr, "error: cannot listen on %s:%d\n", args->host.c_str(), args->port); + return 1; + } + std::fprintf(stderr, "1bit-server %s listening on http://%s:%d (loading %s)\n", kVersion, args->host.c_str(), + args->port, args->model.c_str()); + + std::thread loader([&] { + auto t0 = std::chrono::steady_clock::now(); + auto fail = [&](const std::string& e) { + load_error = e; + std::fprintf(stderr, "error: %s\n", e.c_str()); + state = kFailed; + }; + auto file = GgufFile::open(args->model); + if (!file) return fail(file.error()); + auto tok = Tokenizer::from_gguf(*file); + if (!tok) return fail("tokenizer: " + tok.error()); + auto tmpl = ChatTemplate::from_gguf(*file); + if (!tmpl) return fail(tmpl.error()); + auto model = CpuModel::load(args->model, args->threads, args->ctx_size); + if (!model) return fail(model.error()); + engine.tmpl = std::make_unique(std::move(*tmpl)); + engine.gen = std::make_unique(std::move(*model), std::move(*tok)); + state = kReady; + std::fprintf(stderr, "model ready in %.1f s: %s on cpu, context %u\n", + std::chrono::duration(std::chrono::steady_clock::now() - t0).count(), args->alias.c_str(), + engine.gen->n_ctx()); + }); + loader.detach(); + + return svr.listen_after_bind() ? 0 : 1; +} diff --git a/engine/server/openai.cpp b/engine/server/openai.cpp new file mode 100644 index 00000000..adf05b26 --- /dev/null +++ b/engine/server/openai.cpp @@ -0,0 +1,192 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "openai.h" + +#include +#include + +namespace onebit::openai { + +namespace { + +std::expected parse_common(const json& b, Request& r) { + if (!b.is_object()) return std::unexpected(std::string("request body must be a JSON object")); + auto number = [&](const char* key, auto& field) -> std::expected { + if (!b.contains(key) || b[key].is_null()) return {}; + if (!b[key].is_number()) return std::unexpected(std::format("'{}' must be a number", key)); + field = b[key].get>(); + return {}; + }; + SamplingParams& s = r.gen.sampling; + for (auto res : {number("temperature", s.temperature), number("top_p", s.top_p), number("min_p", s.min_p)}) + if (!res) return res; + if (auto res = number("top_k", s.top_k); !res) return res; + if (b.contains("seed") && b["seed"].is_number_integer() && b["seed"].get() >= 0) { + s.seed = b["seed"].get(); + } else { + s.seed = std::random_device{}(); + } + for (const char* key : {"max_completion_tokens", "max_tokens", "n_predict"}) { + if (b.contains(key) && !b[key].is_null()) { + if (!b[key].is_number_integer()) return std::unexpected(std::format("'{}' must be an integer", key)); + r.gen.max_tokens = b[key].get(); + break; + } + } + if (b.contains("n") && !b["n"].is_null() && b["n"] != 1) + return std::unexpected(std::string("only n = 1 is supported")); + if (b.contains("stop") && !b["stop"].is_null()) { + const json& st = b["stop"]; + if (st.is_string()) { + r.gen.stops.push_back(st.get()); + } else if (st.is_array()) { + for (const auto& x : st) { + if (!x.is_string()) return std::unexpected(std::string("'stop' entries must be strings")); + r.gen.stops.push_back(x.get()); + } + } else { + return std::unexpected(std::string("'stop' must be a string or an array of strings")); + } + } + r.stream = b.value("stream", false); + if (b.contains("stream_options") && b["stream_options"].is_object()) + r.include_usage = b["stream_options"].value("include_usage", false); + return {}; +} + +// Content may be a string, null, or an array of parts; only text parts are +// accepted, joined with newlines. +std::expected text_content(const json& c) { + if (c.is_string() || c.is_null()) return c; + if (!c.is_array()) return std::unexpected(std::string("message content must be a string or an array")); + std::string out; + for (const auto& part : c) { + if (!part.is_object() || part.value("type", "") != "text" || !part.contains("text") || !part["text"].is_string()) + return std::unexpected(std::string("only text content parts are supported")); + if (!out.empty()) out += "\n"; + out += part["text"].get(); + } + return out; +} + +} // namespace + +std::expected parse_chat(const json& b) { + Request r; + r.chat = true; + if (auto res = parse_common(b, r); !res) return std::unexpected(res.error()); + if (!b.contains("messages") || !b["messages"].is_array() || b["messages"].empty()) + return std::unexpected(std::string("'messages' must be a non-empty array")); + r.messages = json::array(); + for (const auto& m : b["messages"]) { + if (!m.is_object() || !m.contains("role") || !m["role"].is_string()) + return std::unexpected(std::string("each message needs a string 'role'")); + json msg = m; + if (m.contains("content")) { + auto c = text_content(m["content"]); + if (!c) return std::unexpected(c.error()); + msg["content"] = *c; + } + r.messages.push_back(std::move(msg)); + } + r.tools = b.contains("tools") && b["tools"].is_array() && !b["tools"].empty() ? b["tools"] : json(nullptr); + r.template_kwargs = b.contains("chat_template_kwargs") && b["chat_template_kwargs"].is_object() + ? b["chat_template_kwargs"] + : json::object(); + r.split_reasoning = b.value("reasoning_format", std::string("auto")) != "none"; + return r; +} + +std::expected parse_completion(const json& b) { + Request r; + r.chat = false; + r.split_reasoning = false; + if (auto res = parse_common(b, r); !res) return std::unexpected(res.error()); + if (!b.contains("prompt")) return std::unexpected(std::string("'prompt' is required")); + const json& p = b["prompt"]; + if (p.is_string()) { + r.prompt = p.get(); + } else if (p.is_array() && p.size() == 1 && p[0].is_string()) { + r.prompt = p[0].get(); + } else { + return std::unexpected(std::string("'prompt' must be a string (batched prompts are not supported)")); + } + return r; +} + +json usage(const GenerateResult& r) { + return {{"prompt_tokens", r.prompt_tokens}, + {"completion_tokens", r.predicted_n}, + {"total_tokens", r.prompt_tokens + r.predicted_n}, + {"prompt_tokens_details", {{"cached_tokens", r.cache_n}}}}; +} + +json timings(const GenerateResult& r) { + const int32_t evaluated = r.prompt_tokens - r.cache_n; + auto per_second = [](int32_t n, double ms) { return ms > 0.0 ? 1000.0 * n / ms : 0.0; }; + return {{"cache_n", r.cache_n}, + {"prompt_n", evaluated}, + {"prompt_ms", r.prompt_ms}, + {"prompt_per_token_ms", evaluated ? r.prompt_ms / evaluated : 0.0}, + {"prompt_per_second", per_second(evaluated, r.prompt_ms)}, + {"predicted_n", r.predicted_n}, + {"predicted_ms", r.predicted_ms}, + {"predicted_per_token_ms", r.predicted_n ? r.predicted_ms / r.predicted_n : 0.0}, + {"predicted_per_second", per_second(r.predicted_n, r.predicted_ms)}}; +} + +json chat_response(const Meta& m, const GenerateResult& r) { + json msg = {{"role", "assistant"}, {"content", r.content}}; + if (!r.reasoning.empty()) msg["reasoning_content"] = r.reasoning; + return {{"id", m.id}, + {"object", "chat.completion"}, + {"created", m.created}, + {"model", m.model}, + {"choices", json::array({{{"index", 0}, {"message", msg}, {"finish_reason", r.finish_reason}}})}, + {"usage", usage(r)}, + {"timings", timings(r)}}; +} + +json chat_chunk(const Meta& m, const json& delta, const json& finish_reason) { + return {{"id", m.id}, + {"object", "chat.completion.chunk"}, + {"created", m.created}, + {"model", m.model}, + {"choices", json::array({{{"index", 0}, {"delta", delta}, {"finish_reason", finish_reason}}})}}; +} + +json completion_response(const Meta& m, const GenerateResult& r) { + return {{"id", m.id}, + {"object", "text_completion"}, + {"created", m.created}, + {"model", m.model}, + {"choices", json::array({{{"index", 0}, {"text", r.content}, {"finish_reason", r.finish_reason}}})}, + {"usage", usage(r)}, + {"timings", timings(r)}}; +} + +json completion_chunk(const Meta& m, const std::string& text, const json& finish_reason) { + return {{"id", m.id}, + {"object", "text_completion"}, + {"created", m.created}, + {"model", m.model}, + {"choices", json::array({{{"index", 0}, {"text", text}, {"finish_reason", finish_reason}}})}}; +} + +json error(int code, const std::string& message, const std::string& type) { + return {{"error", {{"code", code}, {"message", message}, {"type", type}}}}; +} + +} // namespace onebit::openai diff --git a/engine/server/openai.h b/engine/server/openai.h new file mode 100644 index 00000000..790132ce --- /dev/null +++ b/engine/server/openai.h @@ -0,0 +1,67 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// openai.h: OpenAI / llama-server request parsing and response shapes. +// +// Responses carry both OpenAI's `usage` and llama-server's `timings`, which is +// what Lemonade reads for its telemetry. +#pragma once + +#include +#include +#include + +#include + +#include "generator.h" + +namespace onebit::openai { + +using json = nlohmann::ordered_json; + +struct Request { + bool chat = true; + bool stream = false; + bool include_usage = false; // stream_options.include_usage + json messages; // chat: text-only messages + json tools; // chat: array or null + json template_kwargs; // chat: chat_template_kwargs + bool split_reasoning = true; // chat: reasoning_format != "none" + std::string prompt; // completion + GenerateRequest gen; // prompt tokens are filled in by the caller +}; + +// Parse a /v1/chat/completions or /v1/completions body. Unknown fields are +// ignored; fields whose meaning would be silently dropped (n > 1, non-text +// content parts) are rejected. +std::expected parse_chat(const json& body); +std::expected parse_completion(const json& body); + +struct Meta { + std::string id; + std::string model; + int64_t created = 0; +}; + +json usage(const GenerateResult& r); +json timings(const GenerateResult& r); + +json chat_response(const Meta& m, const GenerateResult& r); +json chat_chunk(const Meta& m, const json& delta, const json& finish_reason); +json completion_response(const Meta& m, const GenerateResult& r); +json completion_chunk(const Meta& m, const std::string& text, const json& finish_reason); + +json error(int code, const std::string& message, const std::string& type); + +} // namespace onebit::openai diff --git a/engine/server/text_stream.cpp b/engine/server/text_stream.cpp new file mode 100644 index 00000000..13f378d5 --- /dev/null +++ b/engine/server/text_stream.cpp @@ -0,0 +1,156 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "text_stream.h" + +#include + +namespace onebit { + +namespace { + +constexpr std::string_view kOpen = ""; +constexpr std::string_view kClose = ""; + +bool is_space(char c) { return c == ' ' || c == '\n' || c == '\r' || c == '\t'; } + +// Length of the longest suffix of `s` that is a proper prefix of `word`. +size_t partial_suffix(std::string_view s, std::string_view word) { + for (size_t n = std::min(s.size(), word.size() - 1); n > 0; --n) + if (s.substr(s.size() - n) == word.substr(0, n)) return n; + return 0; +} + +// Bytes at the end of `s` that start an incomplete UTF-8 sequence. +size_t incomplete_utf8_tail(std::string_view s) { + for (size_t back = 1; back <= std::min(3, s.size()); ++back) { + const auto c = static_cast(s[s.size() - back]); + if ((c & 0xC0) == 0x80) continue; // continuation byte: keep looking back + size_t need = (c & 0xE0) == 0xC0 ? 2 : (c & 0xF0) == 0xE0 ? 3 : (c & 0xF8) == 0xF0 ? 4 : 1; + return need > back ? back : 0; + } + return 0; +} + +} // namespace + +TextStream::TextStream(std::vector stops, bool split_reasoning, bool reasoning_open) + : stops_(std::move(stops)), split_(split_reasoning) { + std::erase_if(stops_, [](const std::string& s) { return s.empty(); }); + mode_ = !split_ ? Mode::Content : reasoning_open ? Mode::Reasoning : Mode::Start; +} + +void TextStream::emit(TextDelta& out, std::string_view content, std::string_view reasoning) { + if (!content.empty() && trim_content_start_) { + size_t skip = 0; + while (skip < content.size() && is_space(content[skip])) ++skip; + content.remove_prefix(skip); + if (!content.empty()) trim_content_start_ = false; + } + out.content += content; + out.reasoning += reasoning; + all_content_ += content; + all_reasoning_ += reasoning; +} + +void TextStream::route(std::string_view text, bool final, TextDelta& out) { + route_hold_ += text; + for (;;) { + std::string& h = route_hold_; + if (mode_ == Mode::Content) { + emit(out, h, {}); + h.clear(); + return; + } + if (mode_ == Mode::Start) { + size_t lead = 0; + while (lead < h.size() && is_space(h[lead])) ++lead; + std::string_view rest = std::string_view(h).substr(lead); + if (rest.starts_with(kOpen)) { + h.erase(0, lead + kOpen.size()); + mode_ = Mode::Reasoning; + continue; + } + if (!final && kOpen.starts_with(rest)) return; // may still become + mode_ = Mode::Content; + continue; + } + // Reasoning: emit up to , trimming the block's outer whitespace. + if (all_reasoning_.empty()) { + size_t lead = 0; + while (lead < h.size() && is_space(h[lead])) ++lead; + h.erase(0, lead); + } + if (size_t close = h.find(kClose); close != std::string::npos) { + size_t end = close; + while (end > 0 && is_space(h[end - 1])) --end; + emit(out, {}, std::string_view(h).substr(0, end)); + h.erase(0, close + kClose.size()); + mode_ = Mode::Content; + trim_content_start_ = true; + continue; + } + if (final) { // unterminated block: everything is reasoning + size_t end = h.size(); + while (end > 0 && is_space(h[end - 1])) --end; + emit(out, {}, std::string_view(h).substr(0, end)); + h.clear(); + return; + } + // Hold a possible partial and trailing whitespace (it may precede the tag). + size_t keep = partial_suffix(h, kClose); + size_t end = h.size() - keep; + while (end > 0 && is_space(h[end - 1])) --end; + emit(out, {}, std::string_view(h).substr(0, end)); + h.erase(0, end); + return; + } +} + +TextDelta TextStream::push(std::string_view bytes) { + TextDelta out; + if (stopped_) return out; + + std::string text = utf8_tail_ + std::string(bytes); + const size_t tail = incomplete_utf8_tail(text); + utf8_tail_ = text.substr(text.size() - tail); + text.resize(text.size() - tail); + + stop_hold_ += text; + size_t first = std::string::npos; + for (const auto& s : stops_) first = std::min(first, stop_hold_.find(s)); + if (first != std::string::npos) { + stopped_ = true; + route(std::string_view(stop_hold_).substr(0, first), true, out); + stop_hold_.clear(); + utf8_tail_.clear(); + return out; + } + size_t keep = 0; + for (const auto& s : stops_) keep = std::max(keep, partial_suffix(stop_hold_, s)); + route(std::string_view(stop_hold_).substr(0, stop_hold_.size() - keep), false, out); + stop_hold_.erase(0, stop_hold_.size() - keep); + return out; +} + +TextDelta TextStream::finish() { + TextDelta out; + if (stopped_) return out; + route(stop_hold_, true, out); + stop_hold_.clear(); + utf8_tail_.clear(); + return out; +} + +} // namespace onebit diff --git a/engine/server/text_stream.h b/engine/server/text_stream.h new file mode 100644 index 00000000..5256e425 --- /dev/null +++ b/engine/server/text_stream.h @@ -0,0 +1,74 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// text_stream.h: turns per-token bytes into text deltas that are safe to send. +// +// Three stages, applied in order to the generated bytes: +// 1. UTF-8: a token can end inside a multi-byte character; those bytes wait +// for the rest of the character. +// 2. Stop strings: text that could be the start of a stop string waits until +// it either completes the stop string (generation ends, the stop string is +// not emitted) or cannot. +// 3. Reasoning: with split_reasoning, a leading ... block goes +// to `reasoning` and the rest to `content` (llama-server's reasoning +// format), with the whitespace around the block trimmed. +#pragma once + +#include +#include +#include + +namespace onebit { + +struct TextDelta { + std::string content; + std::string reasoning; + bool empty() const { return content.empty() && reasoning.empty(); } +}; + +class TextStream { +public: + // reasoning_open: the prompt already opened a block, so output + // starts inside it (templates that end the prompt with "\n"). + TextStream(std::vector stops, bool split_reasoning, bool reasoning_open = false); + + // Feeds the bytes of one generated token. Returns what can be sent now. + TextDelta push(std::string_view bytes); + + // True once a stop string has matched; later pushes return nothing. + bool stopped() const { return stopped_; } + + // End of generation: releases held text (an incomplete UTF-8 tail is dropped). + TextDelta finish(); + + // Everything emitted so far. + const std::string& content() const { return all_content_; } + const std::string& reasoning() const { return all_reasoning_; } + +private: + void route(std::string_view text, bool final, TextDelta& out); + void emit(TextDelta& out, std::string_view content, std::string_view reasoning); + + std::vector stops_; + bool split_; + enum class Mode { Start, Reasoning, Content } mode_; + bool trim_content_start_ = false; + bool stopped_ = false; + std::string utf8_tail_; // incomplete UTF-8 sequence + std::string stop_hold_; // text that may begin a stop string + std::string route_hold_; // text held by the reasoning router + std::string all_content_, all_reasoning_; +}; + +} // namespace onebit diff --git a/tests/golden/qwen3-0.6b-capitals/meta.json b/tests/golden/qwen3-0.6b-capitals/meta.json index 121f603b..1b479e4d 100644 --- a/tests/golden/qwen3-0.6b-capitals/meta.json +++ b/tests/golden/qwen3-0.6b-capitals/meta.json @@ -9,5 +9,6 @@ "attn_implementation": "eager", "torch": "2.13.0+cpu", "transformers": "5.13.0.dev0", - "logits_sha256": "3aaec202fa3d5e4dca1ebc053623836a57960f242ccbd34e98036aa97654820a" + "logits_sha256": "3aaec202fa3d5e4dca1ebc053623836a57960f242ccbd34e98036aa97654820a", + "continuation": " Tokyo. The capital of Italy is Rome. The capital of Spain is Madrid. The capital of China is Beijing. The" } \ No newline at end of file diff --git a/tests/golden/qwen3-template/cases.jsonl b/tests/golden/qwen3-template/cases.jsonl new file mode 100644 index 00000000..d6cf3409 --- /dev/null +++ b/tests/golden/qwen3-template/cases.jsonl @@ -0,0 +1,12 @@ +{"tools": null, "add_generation_prompt": true, "kwargs": {}, "messages": [{"role": "user", "content": "Hello!"}], "expected": "<|im_start|>user\nHello!<|im_end|>\n<|im_start|>assistant\n"} +{"tools": null, "add_generation_prompt": true, "kwargs": {"enable_thinking": false}, "messages": [{"role": "user", "content": "Hello!"}], "expected": "<|im_start|>user\nHello!<|im_end|>\n<|im_start|>assistant\n\n\n\n\n"} +{"tools": null, "add_generation_prompt": false, "kwargs": {}, "messages": [{"role": "user", "content": "Hello!"}], "expected": "<|im_start|>user\nHello!<|im_end|>\n"} +{"tools": null, "add_generation_prompt": true, "kwargs": {}, "messages": [{"role": "system", "content": "You are terse."}, {"role": "user", "content": "Hi"}], "expected": "<|im_start|>system\nYou are terse.<|im_end|>\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n"} +{"tools": null, "add_generation_prompt": true, "kwargs": {}, "messages": [{"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "4"}, {"role": "user", "content": "And 3+3?"}], "expected": "<|im_start|>user\nWhat is 2+2?<|im_end|>\n<|im_start|>assistant\n4<|im_end|>\n<|im_start|>user\nAnd 3+3?<|im_end|>\n<|im_start|>assistant\n"} +{"tools": null, "add_generation_prompt": true, "kwargs": {}, "messages": [{"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "\nadd them\n\n\n4"}, {"role": "user", "content": "Thanks"}], "expected": "<|im_start|>user\nWhat is 2+2?<|im_end|>\n<|im_start|>assistant\n4<|im_end|>\n<|im_start|>user\nThanks<|im_end|>\n<|im_start|>assistant\n"} +{"tools": null, "add_generation_prompt": true, "kwargs": {}, "messages": [{"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "4", "reasoning_content": "add them"}, {"role": "user", "content": "Thanks"}], "expected": "<|im_start|>user\nWhat is 2+2?<|im_end|>\n<|im_start|>assistant\n4<|im_end|>\n<|im_start|>user\nThanks<|im_end|>\n<|im_start|>assistant\n"} +{"tools": [{"type": "function", "function": {"name": "get_weather", "description": "Get the current weather for a city.", "parameters": {"type": "object", "properties": {"city": {"type": "string", "description": "City name"}}, "required": ["city"]}}}], "add_generation_prompt": true, "kwargs": {}, "messages": [{"role": "user", "content": "Weather in Paris?"}], "expected": "<|im_start|>system\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n\n{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"description\": \"Get the current weather for a city.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"city\": {\"type\": \"string\", \"description\": \"City name\"}}, \"required\": [\"city\"]}}}\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n<|im_start|>user\nWeather in Paris?<|im_end|>\n<|im_start|>assistant\n"} +{"tools": [{"type": "function", "function": {"name": "get_weather", "description": "Get the current weather for a city.", "parameters": {"type": "object", "properties": {"city": {"type": "string", "description": "City name"}}, "required": ["city"]}}}], "add_generation_prompt": true, "kwargs": {}, "messages": [{"role": "system", "content": "Use tools when useful."}, {"role": "user", "content": "Weather in Paris?"}, {"role": "assistant", "content": "", "tool_calls": [{"type": "function", "function": {"name": "get_weather", "arguments": {"city": "Paris"}}}]}, {"role": "tool", "content": "{\"temp_c\": 18}"}], "expected": "<|im_start|>system\nUse tools when useful.\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n\n{\"type\": \"function\", \"function\": {\"name\": \"get_weather\", \"description\": \"Get the current weather for a city.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"city\": {\"type\": \"string\", \"description\": \"City name\"}}, \"required\": [\"city\"]}}}\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n<|im_start|>user\nWeather in Paris?<|im_end|>\n<|im_start|>assistant\n\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Paris\"}}\n<|im_end|>\n<|im_start|>user\n\n{\"temp_c\": 18}\n<|im_end|>\n<|im_start|>assistant\n"} +{"tools": null, "add_generation_prompt": true, "kwargs": {}, "messages": [{"role": "user", "content": "日本語で答えて 😀 — «quotes» \"escaped\" \\ back"}], "expected": "<|im_start|>user\n日本語で答えて 😀 — «quotes» \"escaped\" \\ back<|im_end|>\n<|im_start|>assistant\n"} +{"tools": null, "add_generation_prompt": true, "kwargs": {}, "messages": [{"role": "user", "content": "line1\nline2\n\n indented\ttab"}], "expected": "<|im_start|>user\nline1\nline2\n\n indented\ttab<|im_end|>\n<|im_start|>assistant\n"} +{"tools": null, "add_generation_prompt": true, "kwargs": {}, "messages": [{"role": "user", "content": ""}], "expected": "<|im_start|>user\n<|im_end|>\n<|im_start|>assistant\n"} diff --git a/tests/golden_template.cpp b/tests/golden_template.cpp new file mode 100644 index 00000000..36aa2e45 --- /dev/null +++ b/tests/golden_template.cpp @@ -0,0 +1,68 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// golden_template: render chat-template cases and compare with HF transformers. +// +// usage: golden_template +// cases.jsonl comes from tools/golden/make_template_golden.py. A case passes +// when the rendered prompt equals HF's byte for byte. +#include +#include +#include + +#include + +#include "chat_template.h" +#include "gguf.h" + +using namespace onebit; +using json = nlohmann::ordered_json; + +int main(int argc, char** argv) { + if (argc != 3) { + std::fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + auto file = GgufFile::open(argv[1]); + if (!file) { + std::fprintf(stderr, "%s\n", file.error().c_str()); + return 1; + } + auto tmpl = ChatTemplate::from_gguf(*file); + if (!tmpl) { + std::fprintf(stderr, "%s\n", tmpl.error().c_str()); + return 1; + } + std::ifstream in(argv[2]); + std::string line; + size_t n = 0, fail = 0; + while (std::getline(in, line)) { + if (line.empty()) continue; + json c = json::parse(line); + ++n; + auto got = tmpl->render(c["messages"], c["tools"], c["add_generation_prompt"].get(), c["kwargs"]); + const std::string want = c["expected"].get(); + if (!got || *got != want) { + ++fail; + std::printf("MISMATCH case %zu\n--- want\n%s\n--- got\n%s\n---\n", n, want.c_str(), + got ? got->c_str() : got.error().c_str()); + } + } + if (n == 0) { + std::fprintf(stderr, "no cases\n"); + return 2; + } + std::printf("%zu cases, mismatches %zu\n%s\n", n, fail, fail ? "FAIL" : "PASS"); + return fail ? 1 : 0; +} diff --git a/tests/server_e2e.py b/tests/server_e2e.py new file mode 100644 index 00000000..79859bae --- /dev/null +++ b/tests/server_e2e.py @@ -0,0 +1,209 @@ +#!/usr/bin/env python3 +# Copyright 2026 bong-water-water-bong +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""End-to-end test of 1bit-server over HTTP, the way Lemonade drives it. + +usage: server_e2e.py --server build/1bit-server --model M.gguf --golden tests/golden/qwen3-0.6b-capitals + +Launches the server with Lemonade's flags, waits on /health, and checks +responses against the golden (HF transformers' greedy continuation) and the +OpenAI / llama-server response contract. Standard library only. +""" +import argparse +import json +import socket +import subprocess +import sys +import time +import urllib.error +import urllib.request + +FAILURES = [] + + +def check(cond, what): + print(("ok " if cond else "FAIL ") + what) + if not cond: + FAILURES.append(what) + + +def free_port(): + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +class Client: + def __init__(self, port): + self.base = "http://127.0.0.1:%d" % port + + def get(self, path): + try: + with urllib.request.urlopen(self.base + path, timeout=30) as r: + return r.status, json.loads(r.read()) + except urllib.error.HTTPError as e: + return e.code, json.loads(e.read() or b"{}") + + def post(self, path, body, raw=None): + data = raw if raw is not None else json.dumps(body).encode() + req = urllib.request.Request(self.base + path, data=data, headers={"Content-Type": "application/json"}) + try: + with urllib.request.urlopen(req, timeout=600) as r: + return r.status, json.loads(r.read()) + except urllib.error.HTTPError as e: + return e.code, json.loads(e.read() or b"{}") + + def stream(self, path, body): + """Returns (list of parsed SSE events, saw [DONE]).""" + req = urllib.request.Request(self.base + path, data=json.dumps(dict(body, stream=True)).encode(), + headers={"Content-Type": "application/json"}) + events, done = [], False + with urllib.request.urlopen(req, timeout=600) as r: + check(r.headers.get("Content-Type", "").startswith("text/event-stream"), "stream content type is SSE") + after_done = 0 + for line in r: + line = line.decode().rstrip("\n") + if not line.startswith("data: "): + continue + payload = line[6:] + if payload == "[DONE]": + done = True + continue + after_done += done + events.append(json.loads(payload)) + check(after_done == 0, "no events after [DONE] (%d events)" % len(events)) + return events, done + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--server", required=True) + ap.add_argument("--model", required=True) + ap.add_argument("--golden", required=True) + args = ap.parse_args() + + meta = json.load(open(args.golden + "/meta.json")) + port = free_port() + # The flags Lemonade's HRX/llama.cpp recipes pass, with our device name. + cmd = [args.server, "-m", args.model, "--ctx-size", "4096", "--device", "cpu", "--port", str(port), + "--jinja", "--metrics", "--parallel", "1"] + proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + c = Client(port) + try: + # /health: 503 while loading, then 200 {"status":"ok"}. + statuses, loading_ok, ready_body = set(), True, None + deadline = time.time() + 300 + while time.time() < deadline: + try: + st, body = c.get("/health") + statuses.add(st) + if st == 200: + ready_body = body + break + loading_ok &= st == 503 and body["error"]["message"] == "Loading model" + except (urllib.error.URLError, ConnectionError): + pass + time.sleep(0.05) + check(loading_ok, "/health is 503 Loading model while loading (seen: %s)" % sorted(statuses)) + check(ready_body == {"status": "ok"}, "/health turns 200 {\"status\": \"ok\"} when ready") + + st, body = c.get("/v1/models") + check(st == 200 and body["object"] == "list" and body["data"][0]["object"] == "model", "/v1/models lists the model") + + # Greedy completion must equal HF transformers' greedy continuation. + req = {"prompt": meta["prompt"], "max_tokens": meta["n_tokens"] - meta["n_prompt"], "temperature": 0} + st, r = c.post("/v1/completions", req) + check(st == 200, "/v1/completions 200") + text = r["choices"][0]["text"] + check(text == meta["continuation"], "greedy completion equals HF continuation: %r" % text) + check(r["choices"][0]["finish_reason"] == "length", "finish_reason length at max_tokens") + check(r["usage"]["prompt_tokens"] == meta["n_prompt"], "usage.prompt_tokens matches the HF tokenization") + check(r["usage"]["completion_tokens"] == req["max_tokens"], "usage.completion_tokens") + t = r["timings"] + check(all(k in t for k in ("prompt_n", "prompt_ms", "predicted_n", "predicted_per_second", "cache_n")), + "timings has the fields Lemonade reads") + + # Same request again: the prompt prefix is reused from the KV cache, same output. + st, r2 = c.post("/v1/completions", req) + check(r2["choices"][0]["text"] == text, "repeat request gives the same text") + check(r2["timings"]["cache_n"] == meta["n_prompt"] - 1, "repeat request reuses the cached prefix (cache_n)") + + # Streaming gives the same text, then a final chunk with timings, then [DONE]. + events, done = c.stream("/v1/completions", dict(req, stream_options={"include_usage": True})) + streamed = "".join(e["choices"][0]["text"] for e in events if e.get("choices")) + check(streamed == text, "streamed completion equals non-streamed") + finals = [e for e in events if e.get("choices") and e["choices"][0]["finish_reason"]] + check(len(finals) == 1 and finals[0]["choices"][0]["finish_reason"] == "length" and "timings" in finals[0], + "one final chunk with finish_reason and timings") + check(events[-1].get("choices") == [] and events[-1]["usage"]["completion_tokens"] == req["max_tokens"], + "include_usage adds a usage chunk") + check(done, "stream ends with [DONE]") + + # Stop strings end generation before the match. + st, r = c.post("/v1/completions", dict(req, stop=[" Rome"])) + check(r["choices"][0]["text"] == meta["continuation"].split(" Rome")[0], "stop string truncates the output") + check(r["choices"][0]["finish_reason"] == "stop", "finish_reason stop on a stop string") + + # Chat, thinking off: plain content, no reasoning. + chat = {"messages": [{"role": "user", "content": "What is the capital of France? One word."}], + "temperature": 0, "max_tokens": 16, "chat_template_kwargs": {"enable_thinking": False}} + st, r = c.post("/v1/chat/completions", chat) + msg = r["choices"][0]["message"] + check(st == 200 and msg["role"] == "assistant", "/v1/chat/completions 200") + check("Paris" in msg["content"] and "" not in msg["content"], "thinking off: answer in content: %r" % msg["content"]) + check("reasoning_content" not in msg, "thinking off: no reasoning_content") + check(r["choices"][0]["finish_reason"] == "stop", "chat ends on the end-of-turn token") + + # Chat, thinking on: the block goes to reasoning_content. + think = dict(chat, max_tokens=512) + del think["chat_template_kwargs"] + st, r = c.post("/v1/chat/completions", think) + msg = r["choices"][0]["message"] + check(msg.get("reasoning_content", "") != "" and "" not in msg["content"], "thinking on: reasoning split out") + check("Paris" in msg["content"], "thinking on: answer in content: %r" % msg["content"][:80]) + events, done = c.stream("/v1/chat/completions", think) + deltas = [e["choices"][0]["delta"] for e in events if e.get("choices")] + check(deltas[0].get("role") == "assistant", "first chat chunk carries the role") + check("".join(d.get("reasoning_content", "") for d in deltas) == msg["reasoning_content"], + "streamed reasoning equals non-streamed") + check("".join(d.get("content") or "" for d in deltas) == msg["content"], "streamed content equals non-streamed") + check(done, "chat stream ends with [DONE]") + + # Errors. + st, r = c.post("/v1/chat/completions", None, raw=b"{not json") + check(st == 400 and r["error"]["type"] == "invalid_request_error", "bad JSON is 400") + st, r = c.post("/v1/chat/completions", dict(chat, n=2)) + check(st == 400, "n=2 is 400") + st, r = c.post("/v1/completions", {"prompt": "word " * 5000, "max_tokens": 1}) + check(st == 400 and r["error"]["type"] == "exceed_context_size_error", "over-long prompt is 400 exceed_context_size_error") + st, r = c.get("/v1/nothing") + check(st == 404, "unknown route is 404") + finally: + proc.terminate() + try: + out = proc.communicate(timeout=10)[0] + except subprocess.TimeoutExpired: + proc.kill() + out = proc.communicate()[0] + if FAILURES: + print("\nserver log:\n" + out) + print("%d FAILED" % len(FAILURES)) + return 1 + print("\nPASS") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_openai.cpp b/tests/test_openai.cpp new file mode 100644 index 00000000..fb8d8a3c --- /dev/null +++ b/tests/test_openai.cpp @@ -0,0 +1,120 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include "check.h" +#include "openai.h" + +using namespace onebit; +using json = openai::json; + +static void test_parse_chat() { + auto r = openai::parse_chat(json::parse(R"({ + "model": "x", "stream": true, "stream_options": {"include_usage": true}, + "temperature": 0, "top_k": 5, "max_tokens": 12, "seed": 7, "stop": ["A", "B"], + "messages": [{"role": "system", "content": "s"}, + {"role": "user", "content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}]}], + "chat_template_kwargs": {"enable_thinking": false}, "reasoning_format": "none", + "unknown_field": 1})")); + REQUIRE(r.has_value()); + CHECK(r->stream && r->include_usage); + CHECK(r->gen.sampling.temperature == 0.0f); + CHECK(r->gen.sampling.top_k == 5); + CHECK(r->gen.sampling.seed == 7); + CHECK(r->gen.max_tokens == 12); + CHECK((r->gen.stops == std::vector{"A", "B"})); + CHECK(r->messages[1]["content"] == "a\nb"); + CHECK(r->template_kwargs["enable_thinking"] == false); + CHECK(!r->split_reasoning); + CHECK(r->tools.is_null()); + + auto bad = [](const char* body) { return !openai::parse_chat(json::parse(body)).has_value(); }; + CHECK(bad(R"({"messages": []})")); + CHECK(bad(R"({"messages": [{"content": "no role"}]})")); + CHECK(bad(R"({"messages": [{"role": "user", "content": "x"}], "n": 2})")); + CHECK(bad(R"({"messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "x"}}]}]})")); + CHECK(bad(R"({"messages": [{"role": "user", "content": "x"}], "temperature": "hot"})")); + CHECK(bad(R"({"messages": [{"role": "user", "content": "x"}], "stop": 5})")); + CHECK(bad(R"([1, 2])")); + // max_completion_tokens wins over max_tokens. + auto both = openai::parse_chat(json::parse( + R"({"messages": [{"role": "user", "content": "x"}], "max_tokens": 5, "max_completion_tokens": 9})")); + CHECK(both && both->gen.max_tokens == 9); +} + +static void test_parse_completion() { + auto r = openai::parse_completion(json::parse(R"({"prompt": "hi", "stop": "\n"})")); + REQUIRE(r.has_value()); + CHECK(r->prompt == "hi" && !r->chat && !r->split_reasoning); + CHECK((r->gen.stops == std::vector{"\n"})); + CHECK(openai::parse_completion(json::parse(R"({"prompt": ["one"]})")).has_value()); + CHECK(!openai::parse_completion(json::parse(R"({"prompt": ["a", "b"]})")).has_value()); + CHECK(!openai::parse_completion(json::parse(R"({"prompt": [1, 2, 3]})")).has_value()); + CHECK(!openai::parse_completion(json::parse(R"({})")).has_value()); +} + +// Structural checks on every response shape (nlohmann brace literals can turn +// an intended object into an array, so the shapes are checked, not eyeballed). +static void test_shapes() { + GenerateResult g; + g.content = "4"; + g.reasoning = "add"; + g.finish_reason = "stop"; + g.prompt_tokens = 10; + g.cache_n = 3; + g.predicted_n = 2; + g.prompt_ms = 70.0; + g.predicted_ms = 40.0; + openai::Meta m{"chatcmpl-1", "model", 123}; + + json c = openai::chat_response(m, g); + CHECK(c.is_object() && c["object"] == "chat.completion"); + REQUIRE(c["choices"].is_array() && c["choices"].size() == 1 && c["choices"][0].is_object()); + CHECK(c["choices"][0]["message"].is_object()); + CHECK(c["choices"][0]["message"]["content"] == "4"); + CHECK(c["choices"][0]["message"]["reasoning_content"] == "add"); + CHECK(c["choices"][0]["finish_reason"] == "stop"); + CHECK(c["usage"].is_object() && c["usage"]["prompt_tokens"] == 10 && c["usage"]["completion_tokens"] == 2); + CHECK(c["usage"]["total_tokens"] == 12); + CHECK(c["usage"]["prompt_tokens_details"].is_object() && c["usage"]["prompt_tokens_details"]["cached_tokens"] == 3); + const json& t = c["timings"]; + CHECK(t.is_object() && t["prompt_n"] == 7 && t["cache_n"] == 3 && t["predicted_n"] == 2); + CHECK(t["prompt_per_second"].get() == 100.0); + CHECK(t["predicted_per_second"].get() == 50.0); + + json ch = openai::chat_chunk(m, json{{"content", "x"}}, nullptr); + REQUIRE(ch["choices"].is_array() && ch["choices"][0].is_object()); + CHECK(ch["object"] == "chat.completion.chunk"); + CHECK(ch["choices"][0]["delta"].is_object() && ch["choices"][0]["delta"]["content"] == "x"); + CHECK(ch["choices"][0]["finish_reason"].is_null()); + json role = openai::chat_chunk(m, json{{"role", "assistant"}, {"content", nullptr}}, nullptr); + CHECK(role["choices"][0]["delta"].is_object() && role["choices"][0]["delta"]["role"] == "assistant"); + json empty = openai::chat_chunk(m, json::object(), "length"); + CHECK(empty["choices"][0]["delta"].is_object() && empty["choices"][0]["delta"].empty()); + + json tc = openai::completion_response(m, g); + REQUIRE(tc["choices"].is_array() && tc["choices"][0].is_object()); + CHECK(tc["object"] == "text_completion" && tc["choices"][0]["text"] == "4"); + json tch = openai::completion_chunk(m, "y", nullptr); + CHECK(tch["choices"][0].is_object() && tch["choices"][0]["text"] == "y"); + + json e = openai::error(400, "bad", "invalid_request_error"); + CHECK(e["error"].is_object() && e["error"]["code"] == 400 && e["error"]["message"] == "bad"); +} + +int main() { + test_parse_chat(); + test_parse_completion(); + test_shapes(); + return test_result("test_openai"); +} diff --git a/tests/test_sampler.cpp b/tests/test_sampler.cpp new file mode 100644 index 00000000..15b35c91 --- /dev/null +++ b/tests/test_sampler.cpp @@ -0,0 +1,75 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include +#include + +#include "check.h" +#include "sampler.h" + +using namespace onebit; + +static std::map histogram(SamplingParams p, const std::vector& logits, int n) { + Sampler s(p); + std::map h; + for (int i = 0; i < n; ++i) ++h[s.sample(logits)]; + return h; +} + +int main() { + const std::vector logits = {1.0f, 3.0f, 2.9f, -5.0f, 0.5f}; + + SamplingParams greedy; + greedy.temperature = 0.0f; + CHECK(Sampler(greedy).sample(logits) == 1); + SamplingParams k1; + k1.top_k = 1; + CHECK(Sampler(k1).sample(logits) == 1); + + // Same seed, same sequence; different seed, (almost surely) different. + SamplingParams hot; + hot.temperature = 1.5f; + hot.top_k = 0; + hot.top_p = 1.0f; + hot.min_p = 0.0f; + hot.seed = 42; + Sampler a(hot), b(hot); + std::vector sa, sb; + for (int i = 0; i < 64; ++i) { + sa.push_back(a.sample(logits)); + sb.push_back(b.sample(logits)); + } + CHECK(sa == sb); + + // top-k 2 only ever yields the two largest logits. + SamplingParams k2 = hot; + k2.top_k = 2; + for (auto [id, n] : histogram(k2, logits, 2000)) CHECK(id == 1 || id == 2); + + // top-p: p(1)~0.49, p(2)~0.44 -> 0.5 needs both; token 0 (~0.05) excluded. + SamplingParams tp = hot; + tp.top_p = 0.5f; + for (auto [id, n] : histogram(tp, logits, 2000)) CHECK(id == 1 || id == 2); + + // min-p 0.5: only candidates with p >= 0.5 * p_max survive (tokens 1 and 2). + SamplingParams mp = hot; + mp.min_p = 0.5f; + for (auto [id, n] : histogram(mp, logits, 2000)) CHECK(id == 1 || id == 2); + + // No filters: every token with non-negligible mass shows up. + auto all = histogram(hot, logits, 20000); + CHECK(all.count(0) && all.count(1) && all.count(2) && all.count(4)); + + return test_result("test_sampler"); +} diff --git a/tests/test_text_stream.cpp b/tests/test_text_stream.cpp new file mode 100644 index 00000000..886d7258 --- /dev/null +++ b/tests/test_text_stream.cpp @@ -0,0 +1,122 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#include +#include + +#include "check.h" +#include "text_stream.h" + +using namespace onebit; + +namespace { + +struct Run { + std::string content, reasoning; + bool stopped = false; +}; + +// Feeds pieces one at a time, as tokens would arrive, and concatenates deltas. +Run feed(const std::vector& pieces, std::vector stops = {}, bool split = false, + bool open = false) { + TextStream ts(std::move(stops), split, open); + Run r; + for (const auto& p : pieces) { + TextDelta d = ts.push(p); + r.content += d.content; + r.reasoning += d.reasoning; + } + TextDelta d = ts.finish(); + r.content += d.content; + r.reasoning += d.reasoning; + r.stopped = ts.stopped(); + CHECK(r.content == ts.content()); + CHECK(r.reasoning == ts.reasoning()); + return r; +} + +// Every way of cutting `s` into single bytes: the strictest streaming case. +std::vector bytes(const std::string& s) { + std::vector out; + for (char c : s) out.emplace_back(1, c); + return out; +} + +} // namespace + +static void test_utf8() { + // "é" split across two tokens must not be emitted half-way. + TextStream ts({}, false); + CHECK(ts.push("caf\xC3").content == "caf"); + CHECK(ts.push("\xA9!").content == "\xC3\xA9!"); + CHECK(feed(bytes("日本語 😀")).content == "日本語 😀"); + // A dangling lead byte at the end is dropped rather than sent broken. + CHECK(feed({"ok\xE4\xBD"}).content == "ok"); +} + +static void test_stops() { + Run r = feed(bytes("Hello world. User: hi"), {"User:"}); + CHECK(r.content == "Hello world. "); + CHECK(r.stopped); + r = feed({"ab", "cUs", "er", "x"}, {"User"}); // stop string spans tokens + CHECK(r.content == "abc"); + CHECK(r.stopped); + r = feed({"Use", "ful"}, {"User"}); // near miss: held, then released + CHECK(r.content == "Useful"); + CHECK(!r.stopped); + r = feed({"Us"}, {"User"}); // ends on a partial stop: released at finish + CHECK(r.content == "Us"); + r = feed({"a", "STOP", "b"}, {"", "STOP"}); // empty stop strings ignored + CHECK(r.content == "a"); + // Earliest of several stop strings wins. + r = feed({"x END y STOP"}, {"STOP", "END"}); + CHECK(r.content == "x "); +} + +static void test_reasoning() { + const std::string text = "\nLet me think.\n\n\nThe answer is 4."; + for (const auto& pieces : {std::vector{text}, bytes(text)}) { + Run r = feed(pieces, {}, true); + CHECK(r.reasoning == "Let me think."); + CHECK(r.content == "The answer is 4."); + } + // Split off: the tags stay in content. + CHECK(feed({text}).content == text); + // No think block: all content, including text that merely starts like ""), {}, true).content == ""); + CHECK(feed({" Hi"}, {}, true).content == " Hi"); + // Prompt already opened the block. + Run r = feed(bytes("step one\n\n\nDone"), {}, true, true); + CHECK(r.reasoning == "step one"); + CHECK(r.content == "Done"); + // Unterminated block (hit max tokens): all of it is reasoning. + r = feed(bytes("\nstill going "), {}, true); + CHECK(r.reasoning == "still going"); + CHECK(r.content.empty()); + // Interior whitespace survives streaming. + r = feed(bytes("a\n\nbc"), {}, true); + CHECK(r.reasoning == "a\n\nb"); + CHECK(r.content == "c"); + // Stop string inside the reasoning ends everything there. + r = feed(bytes("abc STOP defx"), {"STOP"}, true); + CHECK(r.reasoning == "abc"); + CHECK(r.content.empty()); +} + +int main() { + test_utf8(); + test_stops(); + test_reasoning(); + return test_result("test_text_stream"); +} diff --git a/tools/copyright.py b/tools/copyright.py index a8ed9cc7..dd3c811b 100644 --- a/tools/copyright.py +++ b/tools/copyright.py @@ -36,7 +36,7 @@ ] BLOCK_STYLES = [(("*.md",), (""))] -EXEMPT = ["LICENSE", "NOTICE", "*.gguf", "*.json", "*.tsv", "tests/golden/*/tokens.txt"] +EXEMPT = ["LICENSE", "NOTICE", "*.gguf", "*.json", "*.jsonl", "*.tsv", "tests/golden/*/tokens.txt"] def _match(name, patterns): diff --git a/tools/golden/make_golden.py b/tools/golden/make_golden.py index 77c9b9b2..cff159eb 100644 --- a/tools/golden/make_golden.py +++ b/tools/golden/make_golden.py @@ -74,11 +74,11 @@ def main(): "torch": torch.__version__, "transformers": transformers.__version__, "logits_sha256": hashlib.sha256(blob).hexdigest(), + "continuation": tok.decode(gen[len(ids):]), } with open(os.path.join(args.out, "meta.json"), "w") as f: json.dump(meta, f, indent=2) print(json.dumps(meta, indent=2)) - print("continuation:", repr(tok.decode(gen[len(ids):]))) if __name__ == "__main__": diff --git a/tools/golden/make_template_golden.py b/tools/golden/make_template_golden.py new file mode 100644 index 00000000..9f54a38c --- /dev/null +++ b/tools/golden/make_template_golden.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +# Copyright 2026 bong-water-water-bong +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Write chat-template golden cases rendered by HF transformers (jinja2). + +usage: make_template_golden.py --model --out cases.jsonl + +Each line: {"messages", "tools", "add_generation_prompt", "kwargs", "expected"}. +""" +import argparse +import json + +from transformers import AutoTokenizer + +WEATHER_TOOL = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string", "description": "City name"}}, + "required": ["city"], + }, + }, +} + +CASES = [ + {"messages": [{"role": "user", "content": "Hello!"}]}, + {"messages": [{"role": "user", "content": "Hello!"}], "kwargs": {"enable_thinking": False}}, + {"messages": [{"role": "user", "content": "Hello!"}], "add_generation_prompt": False}, + {"messages": [{"role": "system", "content": "You are terse."}, {"role": "user", "content": "Hi"}]}, + {"messages": [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"}, + {"role": "user", "content": "And 3+3?"}, + ]}, + {"messages": [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "\nadd them\n\n\n4"}, + {"role": "user", "content": "Thanks"}, + ]}, + {"messages": [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4", "reasoning_content": "add them"}, + {"role": "user", "content": "Thanks"}, + ]}, + {"messages": [{"role": "user", "content": "Weather in Paris?"}], "tools": [WEATHER_TOOL]}, + {"messages": [ + {"role": "system", "content": "Use tools when useful."}, + {"role": "user", "content": "Weather in Paris?"}, + {"role": "assistant", "content": "", "tool_calls": [ + {"type": "function", "function": {"name": "get_weather", "arguments": {"city": "Paris"}}}]}, + {"role": "tool", "content": "{\"temp_c\": 18}"}, + ], "tools": [WEATHER_TOOL]}, + {"messages": [{"role": "user", "content": "日本語で答えて 😀 — «quotes» \"escaped\" \\ back"}]}, + {"messages": [{"role": "user", "content": "line1\nline2\n\n indented\ttab"}]}, + {"messages": [{"role": "user", "content": ""}]}, +] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--out", required=True) + args = ap.parse_args() + tok = AutoTokenizer.from_pretrained(args.model) + with open(args.out, "w") as f: + for case in CASES: + case = {"tools": None, "add_generation_prompt": True, "kwargs": {}, **case} + case["expected"] = tok.apply_chat_template( + case["messages"], tools=case["tools"], add_generation_prompt=case["add_generation_prompt"], + tokenize=False, **case["kwargs"]) + f.write(json.dumps(case, ensure_ascii=False) + "\n") + print("wrote %d cases to %s" % (len(CASES), args.out)) + + +if __name__ == "__main__": + main()