diff --git a/README.md b/README.md index 38fb555..961761a 100644 --- a/README.md +++ b/README.md @@ -48,9 +48,9 @@ For each question, the server: 4. Keeps only the logits for valid answer codes and applies softmax. 5. Maps those probabilities back to the answer names supplied by the caller. -Questions are evaluated independently. Each question gets its own prompt and -another copy of the state. This prevents one answer from affecting another, -but runtime and input-token use grow with the number and length of questions. +The default MLX and vllm-metal backends evaluate questions independently. The +DiffusionGemma backend instead places all answers in one shared canvas and +runs one denoising forward. See [How it works](docs/how-it-works.md) for the prompt format, answer-code registry, token limits, and confidence calculation. @@ -275,6 +275,11 @@ Select a profile before starting the server: SYSTEM_ONE_MODEL=larger uv run uvicorn system_one_lite.api:app --port 8010 ``` +To serve the same model through vllm-metal, see the +[vllm-metal backend guide](docs/vllm-metal.md). +For seeded, single-forward reads with the 4-bit DiffusionGemma model, see the +[DiffusionGemma backend guide](docs/diffusion-gemma.md). + The demo, benchmark, and evaluation tools also accept `--model default` or `--model larger`. Each supported model has a checked-in registry of answer codes that must remain single tokens with the pinned tokenizer. diff --git a/docs/api.md b/docs/api.md index 4ec1bbf..372c50c 100644 --- a/docs/api.md +++ b/docs/api.md @@ -99,8 +99,8 @@ How likely the statement is to be true. | --- | --- | --- | | `model` | `string` | The engine's own model ID. The default is `mlx-community/Qwen3-1.7B-4bit`. | | `answers` | `map` | One answer per question, keyed by your IDs, in request order. | -| `usage.input_tokens` | `integer` | Total tokens evaluated across the independent question prompts. | -| `usage.output_tokens` | `integer` | Always 0. Nothing is generated. | +| `usage.input_tokens` | `integer` | Prompt tokens evaluated by the selected backend. | +| `usage.output_tokens` | `integer` | 0 for local MLX and read-only DiffusionGemma. The vllm-metal backend uses one internal output token per question for each group of up to 128 options. | ### Choice answer @@ -152,5 +152,5 @@ You get an error, never a wrong answer. Only one inference request runs at a time. If the engine is already in use, the server returns `503` with `detail: "the inference engine is busy"`. -Each question uses a full independent prompt. `usage.input_tokens` counts -the state again for each question. +The default MLX and vllm-metal backends count one full prompt per question. +The DiffusionGemma backend counts one shared prompt for the whole request. diff --git a/docs/diffusion-gemma.md b/docs/diffusion-gemma.md new file mode 100644 index 0000000..6344294 --- /dev/null +++ b/docs/diffusion-gemma.md @@ -0,0 +1,69 @@ +# Run with 4-bit DiffusionGemma + +This backend uses `mlx-community/diffusiongemma-26B-A4B-it-4bit` at the pinned +revision in `engine.py`. It sends every question in one prompt. It places one +active canvas token per question. This fits System One's limit of 64 questions. +One read-only model pass returns the requested answer-code scores. The canvas +uses the same fixed seed for each request. + +The backend needs the `feature/structured-reads` branch of MLX-VLM. That +branch adds `/v1/diffusion/reads` without changing normal generation. + +Start the MLX-VLM server from its checkout: + +```bash +cd /path/to/mlx-vlm +uv run --with-editable . python -m mlx_vlm.server --port 8080 +``` + +Start System One in another terminal: + +```bash +cd server +SYSTEM_ONE_BACKEND=mlx-vlm-diffusion \ + uv run uvicorn system_one_lite.api:app --port 8010 +``` + +Use the compact profile for short, latency-sensitive questions: + +```bash +SYSTEM_ONE_BACKEND=mlx-vlm-diffusion-fast \ + uv run uvicorn system_one_lite.api:app --port 8010 +``` + +The compact profile accepts exactly one question. It removes the long prompt +instructions. Use it only when the state, question, and answer labels make the +task clear without extra guidance. It still runs the full model. Its speed +comes only from the shorter prompt. + +The first System One startup downloads the pinned 4-bit model snapshot. The +weights are about 16.5 GB. Both processes use the same Hugging Face cache. + +The read endpoint accepts prompt tokens, an active seed canvas, and exact token +IDs for each answer slot. It checks every token and position before the model +runs. It scores only the requested answer codes and skips the full vocabulary +head. System One then applies its usual softmax over each question's valid +answer codes. + +This path differs from the vllm-metal backend in two ways: + +- All questions share one prompt and one denoising forward. +- The answer slots can attend to the fixed answer template and to one another. + +The public System One request and response formats do not change. Read-only +canvas work reports zero output tokens because no tokens are committed. + +## Profile the backend + +Start the MLX-VLM server, then run the warmed concurrency sweep from +`server/`: + +```bash +uv run python -m tools.benchmark_diffusion +``` + +Use `--json` to save machine-readable results. Set `--requests`, +`--concurrency`, and `--questions` to change the load. Use +`--profile fast --questions 1` to measure the compact profile. The report +includes request throughput, decision throughput, client latency, and request +round-trip latency. diff --git a/docs/how-it-works.md b/docs/how-it-works.md index 3a060a2..310d92c 100644 --- a/docs/how-it-works.md +++ b/docs/how-it-works.md @@ -1,9 +1,13 @@ # How it works -The engine (`server/src/system_one_lite/engine.py`) never decodes text. It builds one full prompt -for each question. The next token must be an answer code. The engine runs a -forward pass and reads the code probabilities at that position. It generates -no tokens. +The default engine (`server/src/system_one_lite/engine.py`) never decodes text. +It builds one full prompt for each question. The next token must be an answer +code. The engine runs a forward pass and reads the code probabilities at that +position. It generates no tokens. + +The optional DiffusionGemma backend keeps the same public API but uses one +shared prompt and a seeded answer canvas. See [Run with 4-bit +DiffusionGemma](diffusion-gemma.md). This page walks through the pieces so you can read the code alongside it. diff --git a/docs/quickstart.md b/docs/quickstart.md index ea8591e..45cbeb6 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -149,8 +149,8 @@ Things worth noticing: never asks for money back. - These probabilities are model scores, not calibrated odds. A value near one does not prove that an answer is correct. -- `usage.output_tokens` is always 0. The engine never generates. Each question - uses a full independent prompt. +- `usage.output_tokens` is 0 for the default and read-only DiffusionGemma + backends. The default backend uses a full independent prompt per question. ## 4. Use the SDK (optional) diff --git a/docs/vllm-metal.md b/docs/vllm-metal.md new file mode 100644 index 0000000..30e64cb --- /dev/null +++ b/docs/vllm-metal.md @@ -0,0 +1,62 @@ +# vllm-metal backend + +System One can use a separate vllm-metal server instead of loading the model +in the API process. The HTTP API and answer shapes stay the same. + +Install vllm-metal in its own Homebrew environment: + +```bash +brew tap vllm-project/vllm-metal https://github.com/vllm-project/vllm-metal +brew install vllm-project/vllm-metal/vllm-metal +``` + +Start vllm-metal with the same pinned model revision as the answer-code +registry: + +```bash +vllm serve mlx-community/Qwen3-1.7B-4bit \ + --revision 3b1b1768f8f8cf8351c712464f906e86c2b8269e \ + --max-model-len 32769 \ + --port 8000 +``` + +The extra model-length slot is for vllM's internal output token. System One +still limits each prompt to 32,768 input tokens. + +Then start System One in another terminal: + +```bash +cd server +SYSTEM_ONE_BACKEND=vllm-metal \ +SYSTEM_ONE_VLLM_BASE_URL=http://127.0.0.1:8000 \ +uv run uvicorn system_one_lite.api:app --port 8010 +``` + +The backend builds each prompt with System One's pinned tokenizer. It sends +the exact prompt token IDs in completion batches. It also asks vllm-metal +for the exact token IDs of every allowed answer code, then normalizes those +scores at System One's temperature. It does not use a top-k cutoff or guess a +score for a missing answer. Question IDs are still used only to route answers +and are never sent to the model. + +vLLM accepts up to 128 requested token IDs in one read. A question with more +options uses more than one exact read. Each read makes vllm-metal sample one +internal token so it can return the requested log probabilities. The System +One response reports that work in `usage.output_tokens`. The local MLX backend +still reports zero output tokens. + +Before inference, the backend checks both token limits. It also asks +vllm-metal to decode the full answer-code registry. A server tokenizer +mismatch stops the request instead of reading the wrong token IDs. + +## Current limit + +This backend does not use the seeded parallel canvas from vLLM PR 57250. Use +the [4-bit DiffusionGemma backend](diffusion-gemma.md) for that execution +shape. MLX-VLM already implements DiffusionGemma on Apple Silicon, so the +seeded read endpoint lives there instead of adapting vllm-metal's +autoregressive scheduler. + +Keep vllm-metal outside the `server` uv environment. Its release pins an exact +MLX build for its native Metal extension, while System One has its own MLX +dependency. diff --git a/server/src/system_one_lite/api.py b/server/src/system_one_lite/api.py index 5144e80..b3ba23e 100644 --- a/server/src/system_one_lite/api.py +++ b/server/src/system_one_lite/api.py @@ -11,6 +11,7 @@ from fastapi import FastAPI, HTTPException from .engine import Engine, RequestContractError +from .mlx_vlm_diffusion import MlxVlmDiffusionEngine, MlxVlmError from .prompts import as_text, confidence, label from .schemas import ( ChoiceAnswer, @@ -22,11 +23,24 @@ ScoreAnswer, Usage, ) +from .vllm_metal import VllmMetalEngine, VllmMetalError def configured_engine(): """Build the process-wide engine from a named profile or exact model ID.""" - return Engine(os.environ.get("SYSTEM_ONE_MODEL")) + backend = os.environ.get("SYSTEM_ONE_BACKEND", "mlx") + if backend == "mlx": + return Engine(os.environ.get("SYSTEM_ONE_MODEL")) + if backend == "vllm-metal": + return VllmMetalEngine(os.environ.get("SYSTEM_ONE_MODEL")) + if backend == "mlx-vlm-diffusion": + return MlxVlmDiffusionEngine(os.environ.get("SYSTEM_ONE_MODEL")) + if backend == "mlx-vlm-diffusion-fast": + return MlxVlmDiffusionEngine( + os.environ.get("SYSTEM_ONE_MODEL"), + compact=True, + ) + raise ValueError(f"unknown SYSTEM_ONE_BACKEND {backend!r}") def labels_for(question): @@ -86,13 +100,20 @@ def evaluate(req: EvaluateRequest): raise HTTPException(503, "the inference engine is busy") runtime_engine = app.state.engine + engine_questions = [ + (as_text(question.instructions), labels_for(question)) for _, question in items + ] try: results, input_tokens, elapsed_ms = runtime_engine.evaluate( as_text(req.state), - [(as_text(question.instructions), labels_for(question)) for _, question in items], + engine_questions, ) except RequestContractError as error: raise HTTPException(422, str(error)) from error + except VllmMetalError as error: + raise HTTPException(502, str(error)) from error + except MlxVlmError as error: + raise HTTPException(502, str(error)) from error finally: app.state.inference_slot.release() @@ -107,7 +128,10 @@ def evaluate(req: EvaluateRequest): question_id: answer_for(question, probabilities) for (question_id, question), probabilities in zip(items, results) }, - usage=Usage(input_tokens=input_tokens, output_tokens=0), + usage=Usage( + input_tokens=input_tokens, + output_tokens=runtime_engine.output_tokens_for(engine_questions), + ), ) return app diff --git a/server/src/system_one_lite/data/diffusiongemma_26b_a4b_it_4bit_answer_codes.json b/server/src/system_one_lite/data/diffusiongemma_26b_a4b_it_4bit_answer_codes.json new file mode 100644 index 0000000..496ba52 --- /dev/null +++ b/server/src/system_one_lite/data/diffusiongemma_26b_a4b_it_4bit_answer_codes.json @@ -0,0 +1,2898 @@ +{ + "model": "mlx-community/diffusiongemma-26B-A4B-it-4bit", + "model_revision": "a7a81407613811e8ba63af92ac0d852b809e191f", + "tokenizer_sha256": "78d2ea3b75a7968d8d9c5dca2dcbddde6321191ff4c576acd6fe19b957b85be1", + "slot_context": "codes occupy one seeded DiffusionGemma canvas token after a list marker", + "codes": [ + { + "code": "A", + "token": " A", + "token_id": 562 + }, + { + "code": "B", + "token": " B", + "token_id": 603 + }, + { + "code": "C", + "token": " C", + "token_id": 565 + }, + { + "code": "D", + "token": " D", + "token_id": 622 + }, + { + "code": "E", + "token": " E", + "token_id": 645 + }, + { + "code": "F", + "token": " F", + "token_id": 633 + }, + { + "code": "G", + "token": " G", + "token_id": 667 + }, + { + "code": "H", + "token": " H", + "token_id": 640 + }, + { + "code": "I", + "token": " I", + "token_id": 564 + }, + { + "code": "J", + "token": " J", + "token_id": 730 + }, + { + "code": "K", + "token": " K", + "token_id": 751 + }, + { + "code": "L", + "token": " L", + "token_id": 639 + }, + { + "code": "M", + "token": " M", + "token_id": 587 + }, + { + "code": "N", + "token": " N", + "token_id": 646 + }, + { + "code": "O", + "token": " O", + "token_id": 708 + }, + { + "code": "P", + "token": " P", + "token_id": 593 + }, + { + "code": "Q", + "token": " Q", + "token_id": 1366 + }, + { + "code": "R", + "token": " R", + "token_id": 630 + }, + { + "code": "S", + "token": " S", + "token_id": 555 + }, + { + "code": "T", + "token": " T", + "token_id": 558 + }, + { + "code": "U", + "token": " U", + "token_id": 749 + }, + { + "code": "V", + "token": " V", + "token_id": 782 + }, + { + "code": "W", + "token": " W", + "token_id": 649 + }, + { + "code": "X", + "token": " X", + "token_id": 1684 + }, + { + "code": "Y", + "token": " Y", + "token_id": 895 + }, + { + "code": "Z", + "token": " Z", + "token_id": 1491 + }, + { + "code": "AA", + "token": " AA", + "token_id": 23952 + }, + { + "code": "AB", + "token": " AB", + "token_id": 14802 + }, + { + "code": "AC", + "token": " AC", + "token_id": 11659 + }, + { + "code": "AD", + "token": " AD", + "token_id": 7747 + }, + { + "code": "AE", + "token": " AE", + "token_id": 31616 + }, + { + "code": "AF", + "token": " AF", + "token_id": 18026 + }, + { + "code": "AG", + "token": " AG", + "token_id": 17692 + }, + { + "code": "AH", + "token": " AH", + "token_id": 46670 + }, + { + "code": "AI", + "token": " AI", + "token_id": 12498 + }, + { + "code": "AJ", + "token": " AJ", + "token_id": 43155 + }, + { + "code": "AK", + "token": " AK", + "token_id": 31194 + }, + { + "code": "AL", + "token": " AL", + "token_id": 9634 + }, + { + "code": "AM", + "token": " AM", + "token_id": 8151 + }, + { + "code": "AN", + "token": " AN", + "token_id": 7790 + }, + { + "code": "AO", + "token": " AO", + "token_id": 62715 + }, + { + "code": "AP", + "token": " AP", + "token_id": 9943 + }, + { + "code": "AQ", + "token": " AQ", + "token_id": 95915 + }, + { + "code": "AR", + "token": " AR", + "token_id": 8886 + }, + { + "code": "AS", + "token": " AS", + "token_id": 9447 + }, + { + "code": "AT", + "token": " AT", + "token_id": 9928 + }, + { + "code": "AU", + "token": " AU", + "token_id": 43777 + }, + { + "code": "AV", + "token": " AV", + "token_id": 28523 + }, + { + "code": "AW", + "token": " AW", + "token_id": 65292 + }, + { + "code": "AX", + "token": " AX", + "token_id": 74418 + }, + { + "code": "AY", + "token": " AY", + "token_id": 75343 + }, + { + "code": "AZ", + "token": " AZ", + "token_id": 45108 + }, + { + "code": "BA", + "token": " BA", + "token_id": 29619 + }, + { + "code": "BB", + "token": " BB", + "token_id": 23814 + }, + { + "code": "BC", + "token": " BC", + "token_id": 19339 + }, + { + "code": "BD", + "token": " BD", + "token_id": 41669 + }, + { + "code": "BE", + "token": " BE", + "token_id": 11219 + }, + { + "code": "BF", + "token": " BF", + "token_id": 79550 + }, + { + "code": "BG", + "token": " BG", + "token_id": 68637 + }, + { + "code": "BH", + "token": " BH", + "token_id": 51317 + }, + { + "code": "BI", + "token": " BI", + "token_id": 30778 + }, + { + "code": "BJ", + "token": " BJ", + "token_id": 104543 + }, + { + "code": "BK", + "token": " BK", + "token_id": 97121 + }, + { + "code": "BL", + "token": " BL", + "token_id": 15358 + }, + { + "code": "BM", + "token": " BM", + "token_id": 46647 + }, + { + "code": "BN", + "token": " BN", + "token_id": 53260 + }, + { + "code": "BO", + "token": " BO", + "token_id": 13609 + }, + { + "code": "BP", + "token": " BP", + "token_id": 39122 + }, + { + "code": "BR", + "token": " BR", + "token_id": 17494 + }, + { + "code": "BS", + "token": " BS", + "token_id": 33083 + }, + { + "code": "BT", + "token": " BT", + "token_id": 35554 + }, + { + "code": "BU", + "token": " BU", + "token_id": 27129 + }, + { + "code": "BV", + "token": " BV", + "token_id": 74853 + }, + { + "code": "BW", + "token": " BW", + "token_id": 76557 + }, + { + "code": "BX", + "token": " BX", + "token_id": 204893 + }, + { + "code": "BY", + "token": " BY", + "token_id": 13222 + }, + { + "code": "CA", + "token": " CA", + "token_id": 8253 + }, + { + "code": "CB", + "token": " CB", + "token_id": 27331 + }, + { + "code": "CC", + "token": " CC", + "token_id": 10914 + }, + { + "code": "CD", + "token": " CD", + "token_id": 8080 + }, + { + "code": "CE", + "token": " CE", + "token_id": 20200 + }, + { + "code": "CF", + "token": " CF", + "token_id": 24328 + }, + { + "code": "CG", + "token": " CG", + "token_id": 19790 + }, + { + "code": "CH", + "token": " CH", + "token_id": 5831 + }, + { + "code": "CI", + "token": " CI", + "token_id": 29302 + }, + { + "code": "CJ", + "token": " CJ", + "token_id": 56056 + }, + { + "code": "CK", + "token": " CK", + "token_id": 67216 + }, + { + "code": "CL", + "token": " CL", + "token_id": 9837 + }, + { + "code": "CM", + "token": " CM", + "token_id": 20991 + }, + { + "code": "CN", + "token": " CN", + "token_id": 26111 + }, + { + "code": "CO", + "token": " CO", + "token_id": 2671 + }, + { + "code": "CP", + "token": " CP", + "token_id": 23752 + }, + { + "code": "CQ", + "token": " CQ", + "token_id": 77921 + }, + { + "code": "CR", + "token": " CR", + "token_id": 11862 + }, + { + "code": "CS", + "token": " CS", + "token_id": 14923 + }, + { + "code": "CT", + "token": " CT", + "token_id": 21044 + }, + { + "code": "CU", + "token": " CU", + "token_id": 52957 + }, + { + "code": "CV", + "token": " CV", + "token_id": 25572 + }, + { + "code": "CW", + "token": " CW", + "token_id": 62496 + }, + { + "code": "CX", + "token": " CX", + "token_id": 52074 + }, + { + "code": "CY", + "token": " CY", + "token_id": 32931 + }, + { + "code": "CZ", + "token": " CZ", + "token_id": 97354 + }, + { + "code": "DA", + "token": " DA", + "token_id": 21119 + }, + { + "code": "DB", + "token": " DB", + "token_id": 18706 + }, + { + "code": "DC", + "token": " DC", + "token_id": 11829 + }, + { + "code": "DD", + "token": " DD", + "token_id": 38214 + }, + { + "code": "DE", + "token": " DE", + "token_id": 4216 + }, + { + "code": "DF", + "token": " DF", + "token_id": 52894 + }, + { + "code": "DG", + "token": " DG", + "token_id": 61205 + }, + { + "code": "DH", + "token": " DH", + "token_id": 41191 + }, + { + "code": "DI", + "token": " DI", + "token_id": 12271 + }, + { + "code": "DJ", + "token": " DJ", + "token_id": 24089 + }, + { + "code": "DK", + "token": " DK", + "token_id": 74672 + }, + { + "code": "DL", + "token": " DL", + "token_id": 40312 + }, + { + "code": "DM", + "token": " DM", + "token_id": 22529 + }, + { + "code": "DN", + "token": " DN", + "token_id": 50667 + }, + { + "code": "DO", + "token": " DO", + "token_id": 13391 + }, + { + "code": "DP", + "token": " DP", + "token_id": 41385 + }, + { + "code": "DQ", + "token": " DQ", + "token_id": 167505 + }, + { + "code": "DR", + "token": " DR", + "token_id": 19588 + }, + { + "code": "DS", + "token": " DS", + "token_id": 29066 + }, + { + "code": "DT", + "token": " DT", + "token_id": 32528 + }, + { + "code": "DU", + "token": " DU", + "token_id": 41349 + }, + { + "code": "DV", + "token": " DV", + "token_id": 99570 + }, + { + "code": "DW", + "token": " DW", + "token_id": 72185 + }, + { + "code": "DX", + "token": " DX", + "token_id": 70235 + }, + { + "code": "DY", + "token": " DY", + "token_id": 97135 + }, + { + "code": "DZ", + "token": " DZ", + "token_id": 174960 + }, + { + "code": "EA", + "token": " EA", + "token_id": 48005 + }, + { + "code": "EB", + "token": " EB", + "token_id": 72197 + }, + { + "code": "EC", + "token": " EC", + "token_id": 19521 + }, + { + "code": "ED", + "token": " ED", + "token_id": 17026 + }, + { + "code": "EE", + "token": " EE", + "token_id": 27198 + }, + { + "code": "EF", + "token": " EF", + "token_id": 48664 + }, + { + "code": "EG", + "token": " EG", + "token_id": 69228 + }, + { + "code": "EH", + "token": " EH", + "token_id": 113243 + }, + { + "code": "EI", + "token": " EI", + "token_id": 133121 + }, + { + "code": "EJ", + "token": " EJ", + "token_id": 107740 + }, + { + "code": "EK", + "token": " EK", + "token_id": 110764 + }, + { + "code": "EL", + "token": " EL", + "token_id": 17517 + }, + { + "code": "EM", + "token": " EM", + "token_id": 22950 + }, + { + "code": "EN", + "token": " EN", + "token_id": 9163 + }, + { + "code": "EO", + "token": " EO", + "token_id": 128569 + }, + { + "code": "EP", + "token": " EP", + "token_id": 23590 + }, + { + "code": "EQ", + "token": " EQ", + "token_id": 81912 + }, + { + "code": "ER", + "token": " ER", + "token_id": 17360 + }, + { + "code": "ES", + "token": " ES", + "token_id": 21005 + }, + { + "code": "ET", + "token": " ET", + "token_id": 20157 + }, + { + "code": "EU", + "token": " EU", + "token_id": 11514 + }, + { + "code": "EV", + "token": " EV", + "token_id": 15620 + }, + { + "code": "EW", + "token": " EW", + "token_id": 70800 + }, + { + "code": "EX", + "token": " EX", + "token_id": 8495 + }, + { + "code": "EY", + "token": " EY", + "token_id": 99707 + }, + { + "code": "EZ", + "token": " EZ", + "token_id": 105286 + }, + { + "code": "FA", + "token": " FA", + "token_id": 16154 + }, + { + "code": "FB", + "token": " FB", + "token_id": 65834 + }, + { + "code": "FC", + "token": " FC", + "token_id": 21267 + }, + { + "code": "FD", + "token": " FD", + "token_id": 52815 + }, + { + "code": "FE", + "token": " FE", + "token_id": 23185 + }, + { + "code": "FF", + "token": " FF", + "token_id": 50237 + }, + { + "code": "FG", + "token": " FG", + "token_id": 58000 + }, + { + "code": "FH", + "token": " FH", + "token_id": 78021 + }, + { + "code": "FI", + "token": " FI", + "token_id": 16068 + }, + { + "code": "FJ", + "token": " FJ", + "token_id": 185665 + }, + { + "code": "FK", + "token": " FK", + "token_id": 53192 + }, + { + "code": "FL", + "token": " FL", + "token_id": 16586 + }, + { + "code": "FM", + "token": " FM", + "token_id": 32684 + }, + { + "code": "FN", + "token": " FN", + "token_id": 88122 + }, + { + "code": "FO", + "token": " FO", + "token_id": 25880 + }, + { + "code": "FP", + "token": " FP", + "token_id": 40530 + }, + { + "code": "FR", + "token": " FR", + "token_id": 18309 + }, + { + "code": "FS", + "token": " FS", + "token_id": 45925 + }, + { + "code": "FT", + "token": " FT", + "token_id": 37483 + }, + { + "code": "FU", + "token": " FU", + "token_id": 57581 + }, + { + "code": "FV", + "token": " FV", + "token_id": 105235 + }, + { + "code": "FW", + "token": " FW", + "token_id": 91579 + }, + { + "code": "FX", + "token": " FX", + "token_id": 37113 + }, + { + "code": "FY", + "token": " FY", + "token_id": 29692 + }, + { + "code": "GA", + "token": " GA", + "token_id": 23639 + }, + { + "code": "GB", + "token": " GB", + "token_id": 26968 + }, + { + "code": "GC", + "token": " GC", + "token_id": 29013 + }, + { + "code": "GD", + "token": " GD", + "token_id": 46169 + }, + { + "code": "GE", + "token": " GE", + "token_id": 27133 + }, + { + "code": "GF", + "token": " GF", + "token_id": 74659 + }, + { + "code": "GG", + "token": " GG", + "token_id": 83079 + }, + { + "code": "GH", + "token": " GH", + "token_id": 48412 + }, + { + "code": "GI", + "token": " GI", + "token_id": 46577 + }, + { + "code": "GJ", + "token": " GJ", + "token_id": 182529 + }, + { + "code": "GK", + "token": " GK", + "token_id": 126979 + }, + { + "code": "GL", + "token": " GL", + "token_id": 14349 + }, + { + "code": "GM", + "token": " GM", + "token_id": 27506 + }, + { + "code": "GN", + "token": " GN", + "token_id": 85411 + }, + { + "code": "GO", + "token": " GO", + "token_id": 16861 + }, + { + "code": "GP", + "token": " GP", + "token_id": 27379 + }, + { + "code": "GQ", + "token": " GQ", + "token_id": 193930 + }, + { + "code": "GR", + "token": " GR", + "token_id": 17779 + }, + { + "code": "GS", + "token": " GS", + "token_id": 43080 + }, + { + "code": "GT", + "token": " GT", + "token_id": 24179 + }, + { + "code": "GU", + "token": " GU", + "token_id": 37518 + }, + { + "code": "GV", + "token": " GV", + "token_id": 125344 + }, + { + "code": "GW", + "token": " GW", + "token_id": 51667 + }, + { + "code": "GX", + "token": " GX", + "token_id": 154378 + }, + { + "code": "GY", + "token": " GY", + "token_id": 121205 + }, + { + "code": "HA", + "token": " HA", + "token_id": 39487 + }, + { + "code": "HB", + "token": " HB", + "token_id": 53261 + }, + { + "code": "HC", + "token": " HC", + "token_id": 50535 + }, + { + "code": "HD", + "token": " HD", + "token_id": 11876 + }, + { + "code": "HE", + "token": " HE", + "token_id": 13962 + }, + { + "code": "HF", + "token": " HF", + "token_id": 57330 + }, + { + "code": "HG", + "token": " HG", + "token_id": 100326 + }, + { + "code": "HH", + "token": " HH", + "token_id": 65091 + }, + { + "code": "HI", + "token": " HI", + "token_id": 53920 + }, + { + "code": "HJ", + "token": " HJ", + "token_id": 183707 + }, + { + "code": "HK", + "token": " HK", + "token_id": 45353 + }, + { + "code": "HL", + "token": " HL", + "token_id": 61611 + }, + { + "code": "HM", + "token": " HM", + "token_id": 59089 + }, + { + "code": "HN", + "token": " HN", + "token_id": 150602 + }, + { + "code": "HO", + "token": " HO", + "token_id": 30698 + }, + { + "code": "HP", + "token": " HP", + "token_id": 20761 + }, + { + "code": "HQ", + "token": " HQ", + "token_id": 71847 + }, + { + "code": "HR", + "token": " HR", + "token_id": 18341 + }, + { + "code": "HS", + "token": " HS", + "token_id": 32572 + }, + { + "code": "HT", + "token": " HT", + "token_id": 9794 + }, + { + "code": "HU", + "token": " HU", + "token_id": 50835 + }, + { + "code": "HV", + "token": " HV", + "token_id": 52446 + }, + { + "code": "HW", + "token": " HW", + "token_id": 74750 + }, + { + "code": "HX", + "token": " HX", + "token_id": 207837 + }, + { + "code": "HY", + "token": " HY", + "token_id": 41887 + }, + { + "code": "IA", + "token": " IA", + "token_id": 51279 + }, + { + "code": "IB", + "token": " IB", + "token_id": 42143 + }, + { + "code": "IC", + "token": " IC", + "token_id": 19403 + }, + { + "code": "ID", + "token": " ID", + "token_id": 5853 + }, + { + "code": "IE", + "token": " IE", + "token_id": 47780 + }, + { + "code": "IF", + "token": " IF", + "token_id": 14634 + }, + { + "code": "IG", + "token": " IG", + "token_id": 44772 + }, + { + "code": "IH", + "token": " IH", + "token_id": 69657 + }, + { + "code": "II", + "token": " II", + "token_id": 3664 + }, + { + "code": "IJ", + "token": " IJ", + "token_id": 115925 + }, + { + "code": "IK", + "token": " IK", + "token_id": 121834 + }, + { + "code": "IL", + "token": " IL", + "token_id": 15237 + }, + { + "code": "IM", + "token": " IM", + "token_id": 12607 + }, + { + "code": "IN", + "token": " IN", + "token_id": 3247 + }, + { + "code": "IO", + "token": " IO", + "token_id": 13421 + }, + { + "code": "IP", + "token": " IP", + "token_id": 10710 + }, + { + "code": "IQ", + "token": " IQ", + "token_id": 47304 + }, + { + "code": "IR", + "token": " IR", + "token_id": 16459 + }, + { + "code": "IS", + "token": " IS", + "token_id": 5246 + }, + { + "code": "IT", + "token": " IT", + "token_id": 9053 + }, + { + "code": "IU", + "token": " IU", + "token_id": 57150 + }, + { + "code": "IV", + "token": " IV", + "token_id": 10935 + }, + { + "code": "IW", + "token": " IW", + "token_id": 150007 + }, + { + "code": "IX", + "token": " IX", + "token_id": 34838 + }, + { + "code": "IZ", + "token": " IZ", + "token_id": 178335 + }, + { + "code": "JA", + "token": " JA", + "token_id": 48032 + }, + { + "code": "JB", + "token": " JB", + "token_id": 88172 + }, + { + "code": "JC", + "token": " JC", + "token_id": 71592 + }, + { + "code": "JD", + "token": " JD", + "token_id": 57469 + }, + { + "code": "JE", + "token": " JE", + "token_id": 47795 + }, + { + "code": "JF", + "token": " JF", + "token_id": 140634 + }, + { + "code": "JG", + "token": " JG", + "token_id": 164564 + }, + { + "code": "JH", + "token": " JH", + "token_id": 113593 + }, + { + "code": "JI", + "token": " JI", + "token_id": 190023 + }, + { + "code": "JJ", + "token": " JJ", + "token_id": 63669 + }, + { + "code": "JK", + "token": " JK", + "token_id": 112862 + }, + { + "code": "JL", + "token": " JL", + "token_id": 91474 + }, + { + "code": "JM", + "token": " JM", + "token_id": 53191 + }, + { + "code": "JN", + "token": " JN", + "token_id": 236659 + }, + { + "code": "JO", + "token": " JO", + "token_id": 22857 + }, + { + "code": "JP", + "token": " JP", + "token_id": 50603 + }, + { + "code": "JR", + "token": " JR", + "token_id": 56553 + }, + { + "code": "JS", + "token": " JS", + "token_id": 28551 + }, + { + "code": "JT", + "token": " JT", + "token_id": 94258 + }, + { + "code": "JU", + "token": " JU", + "token_id": 85240 + }, + { + "code": "JV", + "token": " JV", + "token_id": 119273 + }, + { + "code": "JW", + "token": " JW", + "token_id": 93191 + }, + { + "code": "JY", + "token": " JY", + "token_id": 234264 + }, + { + "code": "KA", + "token": " KA", + "token_id": 77878 + }, + { + "code": "KB", + "token": " KB", + "token_id": 50633 + }, + { + "code": "KC", + "token": " KC", + "token_id": 39684 + }, + { + "code": "KD", + "token": " KD", + "token_id": 92517 + }, + { + "code": "KE", + "token": " KE", + "token_id": 62574 + }, + { + "code": "KF", + "token": " KF", + "token_id": 150049 + }, + { + "code": "KG", + "token": " KG", + "token_id": 77173 + }, + { + "code": "KH", + "token": " KH", + "token_id": 58894 + }, + { + "code": "KI", + "token": " KI", + "token_id": 76148 + }, + { + "code": "KJ", + "token": " KJ", + "token_id": 161846 + }, + { + "code": "KK", + "token": " KK", + "token_id": 79838 + }, + { + "code": "KL", + "token": " KL", + "token_id": 58228 + }, + { + "code": "KM", + "token": " KM", + "token_id": 48878 + }, + { + "code": "KN", + "token": " KN", + "token_id": 67624 + }, + { + "code": "KO", + "token": " KO", + "token_id": 44479 + }, + { + "code": "KP", + "token": " KP", + "token_id": 74821 + }, + { + "code": "KQ", + "token": " KQ", + "token_id": 121555 + }, + { + "code": "KR", + "token": " KR", + "token_id": 59147 + }, + { + "code": "KS", + "token": " KS", + "token_id": 56122 + }, + { + "code": "KT", + "token": " KT", + "token_id": 89139 + }, + { + "code": "KU", + "token": " KU", + "token_id": 91408 + }, + { + "code": "KV", + "token": " KV", + "token_id": 132182 + }, + { + "code": "KW", + "token": " KW", + "token_id": 109600 + }, + { + "code": "KX", + "token": " KX", + "token_id": 206607 + }, + { + "code": "KY", + "token": " KY", + "token_id": 63346 + }, + { + "code": "KZ", + "token": " KZ", + "token_id": 189327 + }, + { + "code": "LA", + "token": " LA", + "token_id": 11404 + }, + { + "code": "LB", + "token": " LB", + "token_id": 67495 + }, + { + "code": "LC", + "token": " LC", + "token_id": 42110 + }, + { + "code": "LD", + "token": " LD", + "token_id": 36416 + }, + { + "code": "LE", + "token": " LE", + "token_id": 11618 + }, + { + "code": "LF", + "token": " LF", + "token_id": 84769 + }, + { + "code": "LG", + "token": " LG", + "token_id": 30310 + }, + { + "code": "LH", + "token": " LH", + "token_id": 86686 + }, + { + "code": "LI", + "token": " LI", + "token_id": 22979 + }, + { + "code": "LJ", + "token": " LJ", + "token_id": 148316 + }, + { + "code": "LK", + "token": " LK", + "token_id": 144505 + }, + { + "code": "LL", + "token": " LL", + "token_id": 32050 + }, + { + "code": "LM", + "token": " LM", + "token_id": 63457 + }, + { + "code": "LN", + "token": " LN", + "token_id": 112970 + }, + { + "code": "LO", + "token": " LO", + "token_id": 17962 + }, + { + "code": "LP", + "token": " LP", + "token_id": 35849 + }, + { + "code": "LQ", + "token": " LQ", + "token_id": 138139 + }, + { + "code": "LR", + "token": " LR", + "token_id": 45273 + }, + { + "code": "LS", + "token": " LS", + "token_id": 46450 + }, + { + "code": "LT", + "token": " LT", + "token_id": 56635 + }, + { + "code": "LU", + "token": " LU", + "token_id": 79740 + }, + { + "code": "LV", + "token": " LV", + "token_id": 63147 + }, + { + "code": "LW", + "token": " LW", + "token_id": 124651 + }, + { + "code": "LX", + "token": " LX", + "token_id": 82888 + }, + { + "code": "LY", + "token": " LY", + "token_id": 83331 + }, + { + "code": "LZ", + "token": " LZ", + "token_id": 177461 + }, + { + "code": "MA", + "token": " MA", + "token_id": 7883 + }, + { + "code": "MB", + "token": " MB", + "token_id": 18879 + }, + { + "code": "MC", + "token": " MC", + "token_id": 18768 + }, + { + "code": "MD", + "token": " MD", + "token_id": 16781 + }, + { + "code": "ME", + "token": " ME", + "token_id": 15932 + }, + { + "code": "MF", + "token": " MF", + "token_id": 71026 + }, + { + "code": "MG", + "token": " MG", + "token_id": 34815 + }, + { + "code": "MH", + "token": " MH", + "token_id": 60054 + }, + { + "code": "MI", + "token": " MI", + "token_id": 29067 + }, + { + "code": "MJ", + "token": " MJ", + "token_id": 67119 + }, + { + "code": "MK", + "token": " MK", + "token_id": 42372 + }, + { + "code": "ML", + "token": " ML", + "token_id": 20885 + }, + { + "code": "MM", + "token": " MM", + "token_id": 25228 + }, + { + "code": "MN", + "token": " MN", + "token_id": 36261 + }, + { + "code": "MO", + "token": " MO", + "token_id": 17231 + }, + { + "code": "MP", + "token": " MP", + "token_id": 13945 + }, + { + "code": "MQ", + "token": " MQ", + "token_id": 69825 + }, + { + "code": "MR", + "token": " MR", + "token_id": 24291 + }, + { + "code": "MS", + "token": " MS", + "token_id": 12501 + }, + { + "code": "MT", + "token": " MT", + "token_id": 27511 + }, + { + "code": "MU", + "token": " MU", + "token_id": 46490 + }, + { + "code": "MV", + "token": " MV", + "token_id": 53773 + }, + { + "code": "MW", + "token": " MW", + "token_id": 35920 + }, + { + "code": "MX", + "token": " MX", + "token_id": 45796 + }, + { + "code": "MY", + "token": " MY", + "token_id": 27714 + }, + { + "code": "MZ", + "token": " MZ", + "token_id": 154716 + }, + { + "code": "NA", + "token": " NA", + "token_id": 19674 + }, + { + "code": "NB", + "token": " NB", + "token_id": 55740 + }, + { + "code": "NC", + "token": " NC", + "token_id": 23316 + }, + { + "code": "ND", + "token": " ND", + "token_id": 36437 + }, + { + "code": "NE", + "token": " NE", + "token_id": 8874 + }, + { + "code": "NF", + "token": " NF", + "token_id": 42167 + }, + { + "code": "NG", + "token": " NG", + "token_id": 65799 + }, + { + "code": "NH", + "token": " NH", + "token_id": 22869 + }, + { + "code": "NI", + "token": " NI", + "token_id": 32426 + }, + { + "code": "NJ", + "token": " NJ", + "token_id": 31566 + }, + { + "code": "NK", + "token": " NK", + "token_id": 46314 + }, + { + "code": "NL", + "token": " NL", + "token_id": 51900 + }, + { + "code": "NM", + "token": " NM", + "token_id": 44285 + }, + { + "code": "NN", + "token": " NN", + "token_id": 60800 + }, + { + "code": "NO", + "token": " NO", + "token_id": 9424 + }, + { + "code": "NP", + "token": " NP", + "token_id": 40098 + }, + { + "code": "NR", + "token": " NR", + "token_id": 49428 + }, + { + "code": "NS", + "token": " NS", + "token_id": 10707 + }, + { + "code": "NT", + "token": " NT", + "token_id": 24505 + }, + { + "code": "NU", + "token": " NU", + "token_id": 93241 + }, + { + "code": "NV", + "token": " NV", + "token_id": 57484 + }, + { + "code": "NW", + "token": " NW", + "token_id": 45753 + }, + { + "code": "NX", + "token": " NX", + "token_id": 137707 + }, + { + "code": "NY", + "token": " NY", + "token_id": 13476 + }, + { + "code": "NZ", + "token": " NZ", + "token_id": 55990 + }, + { + "code": "OA", + "token": " OA", + "token_id": 76762 + }, + { + "code": "OB", + "token": " OB", + "token_id": 32610 + }, + { + "code": "OC", + "token": " OC", + "token_id": 58646 + }, + { + "code": "OD", + "token": " OD", + "token_id": 43680 + }, + { + "code": "OE", + "token": " OE", + "token_id": 115092 + }, + { + "code": "OF", + "token": " OF", + "token_id": 3639 + }, + { + "code": "OG", + "token": " OG", + "token_id": 75702 + }, + { + "code": "OH", + "token": " OH", + "token_id": 28963 + }, + { + "code": "OI", + "token": " OI", + "token_id": 116687 + }, + { + "code": "OJ", + "token": " OJ", + "token_id": 212412 + }, + { + "code": "OK", + "token": " OK", + "token_id": 16119 + }, + { + "code": "OL", + "token": " OL", + "token_id": 53418 + }, + { + "code": "OM", + "token": " OM", + "token_id": 69261 + }, + { + "code": "ON", + "token": " ON", + "token_id": 8203 + }, + { + "code": "OO", + "token": " OO", + "token_id": 129623 + }, + { + "code": "OP", + "token": " OP", + "token_id": 19577 + }, + { + "code": "OR", + "token": " OR", + "token_id": 8005 + }, + { + "code": "OS", + "token": " OS", + "token_id": 16080 + }, + { + "code": "OT", + "token": " OT", + "token_id": 53341 + }, + { + "code": "OU", + "token": " OU", + "token_id": 103247 + }, + { + "code": "OV", + "token": " OV", + "token_id": 128433 + }, + { + "code": "OW", + "token": " OW", + "token_id": 51986 + }, + { + "code": "OX", + "token": " OX", + "token_id": 72876 + }, + { + "code": "OZ", + "token": " OZ", + "token_id": 210259 + }, + { + "code": "PA", + "token": " PA", + "token_id": 13187 + }, + { + "code": "PB", + "token": " PB", + "token_id": 40920 + }, + { + "code": "PC", + "token": " PC", + "token_id": 7294 + }, + { + "code": "PD", + "token": " PD", + "token_id": 24313 + }, + { + "code": "PE", + "token": " PE", + "token_id": 20375 + }, + { + "code": "PF", + "token": " PF", + "token_id": 53839 + }, + { + "code": "PG", + "token": " PG", + "token_id": 43447 + }, + { + "code": "PH", + "token": " PH", + "token_id": 12462 + }, + { + "code": "PI", + "token": " PI", + "token_id": 6108 + }, + { + "code": "PJ", + "token": " PJ", + "token_id": 95080 + }, + { + "code": "PK", + "token": " PK", + "token_id": 32743 + }, + { + "code": "PL", + "token": " PL", + "token_id": 9123 + }, + { + "code": "PM", + "token": " PM", + "token_id": 8711 + }, + { + "code": "PN", + "token": " PN", + "token_id": 55561 + }, + { + "code": "PO", + "token": " PO", + "token_id": 13174 + }, + { + "code": "PP", + "token": " PP", + "token_id": 27115 + }, + { + "code": "PQ", + "token": " PQ", + "token_id": 106694 + }, + { + "code": "PR", + "token": " PR", + "token_id": 9206 + }, + { + "code": "PS", + "token": " PS", + "token_id": 11156 + }, + { + "code": "PT", + "token": " PT", + "token_id": 26912 + }, + { + "code": "PU", + "token": " PU", + "token_id": 25358 + }, + { + "code": "PV", + "token": " PV", + "token_id": 32985 + }, + { + "code": "PW", + "token": " PW", + "token_id": 50933 + }, + { + "code": "PX", + "token": " PX", + "token_id": 127287 + }, + { + "code": "PY", + "token": " PY", + "token_id": 70643 + }, + { + "code": "PZ", + "token": " PZ", + "token_id": 151048 + }, + { + "code": "QA", + "token": " QA", + "token_id": 73088 + }, + { + "code": "QB", + "token": " QB", + "token_id": 81378 + }, + { + "code": "QC", + "token": " QC", + "token_id": 66419 + }, + { + "code": "QD", + "token": " QD", + "token_id": 131260 + }, + { + "code": "QE", + "token": " QE", + "token_id": 179889 + }, + { + "code": "QH", + "token": " QH", + "token_id": 228576 + }, + { + "code": "QI", + "token": " QI", + "token_id": 215904 + }, + { + "code": "QM", + "token": " QM", + "token_id": 236116 + }, + { + "code": "QP", + "token": " QP", + "token_id": 176850 + }, + { + "code": "QQ", + "token": " QQ", + "token_id": 179904 + }, + { + "code": "QR", + "token": " QR", + "token_id": 44426 + }, + { + "code": "QS", + "token": " QS", + "token_id": 111800 + }, + { + "code": "QT", + "token": " QT", + "token_id": 132413 + }, + { + "code": "QU", + "token": " QU", + "token_id": 29809 + }, + { + "code": "QV", + "token": " QV", + "token_id": 190489 + }, + { + "code": "RA", + "token": " RA", + "token_id": 21242 + }, + { + "code": "RB", + "token": " RB", + "token_id": 44282 + }, + { + "code": "RC", + "token": " RC", + "token_id": 34023 + }, + { + "code": "RD", + "token": " RD", + "token_id": 53493 + }, + { + "code": "RE", + "token": " RE", + "token_id": 4769 + }, + { + "code": "RF", + "token": " RF", + "token_id": 28049 + }, + { + "code": "RG", + "token": " RG", + "token_id": 71373 + }, + { + "code": "RH", + "token": " RH", + "token_id": 55360 + }, + { + "code": "RI", + "token": " RI", + "token_id": 51720 + }, + { + "code": "RJ", + "token": " RJ", + "token_id": 69430 + }, + { + "code": "RK", + "token": " RK", + "token_id": 93622 + }, + { + "code": "RL", + "token": " RL", + "token_id": 61554 + }, + { + "code": "RM", + "token": " RM", + "token_id": 25543 + }, + { + "code": "RN", + "token": " RN", + "token_id": 47457 + }, + { + "code": "RO", + "token": " RO", + "token_id": 17006 + }, + { + "code": "RP", + "token": " RP", + "token_id": 48269 + }, + { + "code": "RQ", + "token": " RQ", + "token_id": 215785 + }, + { + "code": "RR", + "token": " RR", + "token_id": 50013 + }, + { + "code": "RS", + "token": " RS", + "token_id": 20949 + }, + { + "code": "RT", + "token": " RT", + "token_id": 23094 + }, + { + "code": "RU", + "token": " RU", + "token_id": 108257 + }, + { + "code": "RV", + "token": " RV", + "token_id": 36662 + }, + { + "code": "RW", + "token": " RW", + "token_id": 75067 + }, + { + "code": "RX", + "token": " RX", + "token_id": 39625 + }, + { + "code": "RY", + "token": " RY", + "token_id": 147326 + }, + { + "code": "SA", + "token": " SA", + "token_id": 15850 + }, + { + "code": "SB", + "token": " SB", + "token_id": 37122 + }, + { + "code": "SC", + "token": " SC", + "token_id": 15009 + }, + { + "code": "SD", + "token": " SD", + "token_id": 15336 + }, + { + "code": "SE", + "token": " SE", + "token_id": 7952 + }, + { + "code": "SF", + "token": " SF", + "token_id": 33152 + }, + { + "code": "SG", + "token": " SG", + "token_id": 46538 + }, + { + "code": "SH", + "token": " SH", + "token_id": 9586 + }, + { + "code": "SI", + "token": " SI", + "token_id": 24732 + }, + { + "code": "SJ", + "token": " SJ", + "token_id": 75248 + }, + { + "code": "SK", + "token": " SK", + "token_id": 22246 + }, + { + "code": "SL", + "token": " SL", + "token_id": 17700 + }, + { + "code": "SM", + "token": " SM", + "token_id": 15233 + }, + { + "code": "SN", + "token": " SN", + "token_id": 19818 + }, + { + "code": "SO", + "token": " SO", + "token_id": 12583 + }, + { + "code": "SP", + "token": " SP", + "token_id": 11446 + }, + { + "code": "SQ", + "token": " SQ", + "token_id": 55658 + }, + { + "code": "SR", + "token": " SR", + "token_id": 28155 + }, + { + "code": "SS", + "token": " SS", + "token_id": 18236 + }, + { + "code": "ST", + "token": " ST", + "token_id": 4439 + }, + { + "code": "SU", + "token": " SU", + "token_id": 15625 + }, + { + "code": "SV", + "token": " SV", + "token_id": 34854 + }, + { + "code": "SW", + "token": " SW", + "token_id": 19175 + }, + { + "code": "SX", + "token": " SX", + "token_id": 87863 + }, + { + "code": "SY", + "token": " SY", + "token_id": 20288 + }, + { + "code": "SZ", + "token": " SZ", + "token_id": 112930 + }, + { + "code": "TA", + "token": " TA", + "token_id": 39337 + }, + { + "code": "TB", + "token": " TB", + "token_id": 34949 + }, + { + "code": "TC", + "token": " TC", + "token_id": 41393 + }, + { + "code": "TD", + "token": " TD", + "token_id": 38952 + }, + { + "code": "TE", + "token": " TE", + "token_id": 17423 + }, + { + "code": "TF", + "token": " TF", + "token_id": 56165 + }, + { + "code": "TG", + "token": " TG", + "token_id": 72268 + }, + { + "code": "TH", + "token": " TH", + "token_id": 6925 + }, + { + "code": "TI", + "token": " TI", + "token_id": 43988 + }, + { + "code": "TJ", + "token": " TJ", + "token_id": 88129 + }, + { + "code": "TK", + "token": " TK", + "token_id": 63814 + }, + { + "code": "TL", + "token": " TL", + "token_id": 34147 + }, + { + "code": "TM", + "token": " TM", + "token_id": 38610 + }, + { + "code": "TN", + "token": " TN", + "token_id": 26493 + }, + { + "code": "TO", + "token": " TO", + "token_id": 6760 + }, + { + "code": "TP", + "token": " TP", + "token_id": 34352 + }, + { + "code": "TR", + "token": " TR", + "token_id": 8642 + }, + { + "code": "TS", + "token": " TS", + "token_id": 34654 + }, + { + "code": "TT", + "token": " TT", + "token_id": 40774 + }, + { + "code": "TU", + "token": " TU", + "token_id": 52263 + }, + { + "code": "TV", + "token": " TV", + "token_id": 6115 + }, + { + "code": "TW", + "token": " TW", + "token_id": 48704 + }, + { + "code": "TX", + "token": " TX", + "token_id": 27608 + }, + { + "code": "TY", + "token": " TY", + "token_id": 60057 + }, + { + "code": "TZ", + "token": " TZ", + "token_id": 216288 + }, + { + "code": "UA", + "token": " UA", + "token_id": 84168 + }, + { + "code": "UB", + "token": " UB", + "token_id": 90890 + }, + { + "code": "UC", + "token": " UC", + "token_id": 50606 + }, + { + "code": "UD", + "token": " UD", + "token_id": 74765 + }, + { + "code": "UE", + "token": " UE", + "token_id": 32646 + }, + { + "code": "UF", + "token": " UF", + "token_id": 93592 + }, + { + "code": "UG", + "token": " UG", + "token_id": 151515 + }, + { + "code": "UH", + "token": " UH", + "token_id": 134553 + }, + { + "code": "UI", + "token": " UI", + "token_id": 9711 + }, + { + "code": "UK", + "token": " UK", + "token_id": 6322 + }, + { + "code": "UL", + "token": " UL", + "token_id": 56997 + }, + { + "code": "UM", + "token": " UM", + "token_id": 67201 + }, + { + "code": "UN", + "token": " UN", + "token_id": 7866 + }, + { + "code": "UP", + "token": " UP", + "token_id": 24252 + }, + { + "code": "UR", + "token": " UR", + "token_id": 46412 + }, + { + "code": "US", + "token": " US", + "token_id": 2590 + }, + { + "code": "UT", + "token": " UT", + "token_id": 29297 + }, + { + "code": "UU", + "token": " UU", + "token_id": 109410 + }, + { + "code": "UV", + "token": " UV", + "token_id": 19296 + }, + { + "code": "UW", + "token": " UW", + "token_id": 57623 + }, + { + "code": "UX", + "token": " UX", + "token_id": 69873 + }, + { + "code": "VA", + "token": " VA", + "token_id": 30578 + }, + { + "code": "VB", + "token": " VB", + "token_id": 101537 + }, + { + "code": "VC", + "token": " VC", + "token_id": 51191 + }, + { + "code": "VD", + "token": " VD", + "token_id": 158142 + }, + { + "code": "VE", + "token": " VE", + "token_id": 33697 + }, + { + "code": "VF", + "token": " VF", + "token_id": 138963 + }, + { + "code": "VG", + "token": " VG", + "token_id": 95218 + }, + { + "code": "VH", + "token": " VH", + "token_id": 164123 + }, + { + "code": "VI", + "token": " VI", + "token_id": 18603 + }, + { + "code": "VK", + "token": " VK", + "token_id": 54609 + }, + { + "code": "VL", + "token": " VL", + "token_id": 86868 + }, + { + "code": "VM", + "token": " VM", + "token_id": 37231 + }, + { + "code": "VN", + "token": " VN", + "token_id": 90326 + }, + { + "code": "VO", + "token": " VO", + "token_id": 59054 + }, + { + "code": "VP", + "token": " VP", + "token_id": 42752 + }, + { + "code": "VR", + "token": " VR", + "token_id": 28644 + }, + { + "code": "VS", + "token": " VS", + "token_id": 38097 + }, + { + "code": "VT", + "token": " VT", + "token_id": 55503 + }, + { + "code": "VU", + "token": " VU", + "token_id": 171279 + }, + { + "code": "VV", + "token": " VV", + "token_id": 151030 + }, + { + "code": "VW", + "token": " VW", + "token_id": 70365 + }, + { + "code": "VX", + "token": " VX", + "token_id": 113564 + }, + { + "code": "VY", + "token": " VY", + "token_id": 220191 + }, + { + "code": "WA", + "token": " WA", + "token_id": 31582 + }, + { + "code": "WB", + "token": " WB", + "token_id": 72773 + }, + { + "code": "WC", + "token": " WC", + "token_id": 48380 + }, + { + "code": "WD", + "token": " WD", + "token_id": 100311 + }, + { + "code": "WE", + "token": " WE", + "token_id": 19716 + }, + { + "code": "WF", + "token": " WF", + "token_id": 113026 + }, + { + "code": "WG", + "token": " WG", + "token_id": 121848 + }, + { + "code": "WH", + "token": " WH", + "token_id": 16169 + }, + { + "code": "WI", + "token": " WI", + "token_id": 55927 + }, + { + "code": "WJ", + "token": " WJ", + "token_id": 193516 + }, + { + "code": "WK", + "token": " WK", + "token_id": 125879 + }, + { + "code": "WL", + "token": " WL", + "token_id": 80760 + }, + { + "code": "WM", + "token": " WM", + "token_id": 42668 + }, + { + "code": "WO", + "token": " WO", + "token_id": 42626 + } + ] +} diff --git a/server/src/system_one_lite/engine.py b/server/src/system_one_lite/engine.py index 8646d93..6c1abce 100644 --- a/server/src/system_one_lite/engine.py +++ b/server/src/system_one_lite/engine.py @@ -23,13 +23,19 @@ QWEN3_1_7B_MODEL = "mlx-community/Qwen3-1.7B-4bit" QWEN3_4B_MODEL = "mlx-community/Qwen3-4B-Instruct-2507-4bit" +DIFFUSION_GEMMA_MODEL = "mlx-community/diffusiongemma-26B-A4B-it-4bit" MODEL_REVISIONS = { QWEN3_1_7B_MODEL: "3b1b1768f8f8cf8351c712464f906e86c2b8269e", QWEN3_4B_MODEL: "50d427756c6b1b2fe0c0a10f67fbda1fc8e82c1b", + DIFFUSION_GEMMA_MODEL: "a7a81407613811e8ba63af92ac0d852b809e191f", } DEFAULT_MODEL = QWEN3_1_7B_MODEL LARGER_MODEL = QWEN3_4B_MODEL -MODEL_PROFILES = {"default": DEFAULT_MODEL, "larger": LARGER_MODEL} +MODEL_PROFILES = { + "default": DEFAULT_MODEL, + "larger": LARGER_MODEL, + "diffusion": DIFFUSION_GEMMA_MODEL, +} CODES_FILE = files("system_one_lite.data").joinpath("qwen3_1_7b_4bit_answer_codes.json") TEMPERATURE = 0.7 MAX_PROMPT_TOKENS = 32_768 @@ -47,6 +53,7 @@ ANSWER_CODE_CANDIDATES = tuple(string.ascii_uppercase) + tuple( a + b for a in string.ascii_uppercase for b in string.ascii_uppercase ) +MAX_ANSWER_CODES = 578 _REGISTERED_IDS = object() @@ -129,7 +136,25 @@ def answer_code_entries(tokenizer): ) if len(cuts) != 1: raise ValueError(f"codes land at different cut positions: {sorted(cuts)}") - return entries + return entries[:MAX_ANSWER_CODES] + + +def diffusion_answer_code_entries(tokenizer): + """Build codes that occupy one seeded-canvas token after a list marker.""" + entries = [] + for code in ANSWER_CODE_CANDIDATES: + token_ids = tokenizer.encode(" " + code, add_special_tokens=False) + if len(token_ids) != 1: + continue + token_id = token_ids[0] + entries.append( + { + "code": code, + "token": tokenizer.decode([token_id]), + "token_id": token_id, + } + ) + return entries[:MAX_ANSWER_CODES] def _registry_files(): @@ -143,8 +168,8 @@ def _registry_files(): ) -def load_code_registry(model_id, tokenizer, model_path, model_revision): - """Find and fully validate one model's answer-code registry.""" +def packaged_code_registry(model_id): + """Load one packaged registry without loading the model or tokenizer.""" matches = [] for path in _registry_files(): data = json.loads(path.read_text()) @@ -157,6 +182,19 @@ def load_code_registry(model_id, tokenizer, model_path, model_revision): ) path, data = matches[0] + entries = data.get("codes") or [] + codes = tuple(entry["code"] for entry in entries) + token_ids = tuple(entry["token_id"] for entry in entries) + if not codes or codes[0] != "A": + raise ValueError(f"{path.name} must start with the A answer code") + if len(codes) != len(set(codes)) or len(token_ids) != len(set(token_ids)): + raise ValueError(f"{path.name} contains duplicate codes or token ids") + return path, data, codes, token_ids + + +def load_code_registry(model_id, tokenizer, model_path, model_revision): + """Find and fully validate one model's answer-code registry.""" + path, data, codes, token_ids = packaged_code_registry(model_id) expected_hash = tokenizer_sha256(model_path) if data.get("model_revision") != model_revision: raise ValueError( @@ -166,16 +204,14 @@ def load_code_registry(model_id, tokenizer, model_path, model_revision): if data.get("tokenizer_sha256") != expected_hash: raise ValueError(f"{path.name} does not match the loaded tokenizer; regenerate it") - expected = answer_code_entries(tokenizer) + if model_id == DIFFUSION_GEMMA_MODEL: + expected = diffusion_answer_code_entries(tokenizer) + else: + expected = answer_code_entries(tokenizer) if data.get("codes") != expected: raise ValueError(f"{path.name} is incomplete or invalid for {model_id}; regenerate it") - codes = tuple(entry["code"] for entry in expected) - token_ids = tuple(entry["token_id"] for entry in expected) - if not codes or codes[0] != "A": - raise ValueError(f"{path.name} must start with the A answer code") - if len(codes) != len(set(codes)) or len(token_ids) != len(set(token_ids)): - raise ValueError(f"{path.name} contains duplicate codes or token ids") - validate_code_contexts(tokenizer, codes, token_ids) + if model_id != DIFFUSION_GEMMA_MODEL: + validate_code_contexts(tokenizer, codes, token_ids) return path, codes, token_ids @@ -244,6 +280,10 @@ def __init__(self, model_id=None): # warm up: compile Metal kernels before the first real request self.evaluate("warm up", [("pick one", ["yes", "no"])]) + def output_tokens_for(self, questions): + """Return the generated-token count for one prepared request.""" + return 0 + def evaluate(self, state, questions, template=None): """Run one independent full prompt per question. diff --git a/server/src/system_one_lite/mlx_vlm_diffusion.py b/server/src/system_one_lite/mlx_vlm_diffusion.py new file mode 100644 index 0000000..ff8517e --- /dev/null +++ b/server/src/system_one_lite/mlx_vlm_diffusion.py @@ -0,0 +1,263 @@ +"""System One backend for seeded DiffusionGemma reads through MLX-VLM.""" + +import json +import math +import os +import random +import time +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +from mlx_lm.utils import load_tokenizer + +from .engine import ( + DIFFUSION_GEMMA_MODEL, + MAX_PROMPT_TOKENS, + MAX_TOTAL_INPUT_TOKENS, + TEMPERATURE, + RequestContractError, + TooManyOptions, + diffusion_answer_code_entries, + packaged_code_registry, + resolve_model_id, + resolve_model_snapshot, + tokenizer_sha256, +) + +DEFAULT_BASE_URL = "http://127.0.0.1:8080" +DEFAULT_TIMEOUT = 120.0 +DEFAULT_DIFFUSION_SEED = 42 + + +class MlxVlmError(RuntimeError): + """The MLX-VLM server could not complete a structured diffusion read.""" + + +def restricted_softmax(logprobs, temperature=TEMPERATURE): + """Renormalize vocabulary log probabilities over allowed answer codes.""" + scaled = [value / temperature for value in logprobs] + peak = max(scaled) + weights = [math.exp(value - peak) for value in scaled] + total = sum(weights) + return [value / total for value in weights] + + +def _encode(tokenizer, text): + return tokenizer.encode(text, add_special_tokens=False) + + +def build_prompt(tokenizer, state, questions, codes, compact=False): + """Build one encoder prompt containing every independent question.""" + if compact: + parts = [state.strip()] + for index, (instructions, labels) in enumerate(questions, 1): + prefix = "" if len(questions) == 1 else f"q{index} " + options = " ".join(f"{codes[i]} {item}" for i, item in enumerate(labels)) + parts.append(f" {prefix}{instructions} {options}") + text = tokenizer.apply_chat_template( + [{"role": "user", "content": "".join(parts)}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + return _encode(tokenizer, text) + + parts = [ + "Answer each question about the state. Each question lists its allowed answers.\n\n", + "State:\n", + state.strip(), + ] + for index, (instructions, labels) in enumerate(questions, 1): + parts.append(f"\n\nQuestion q{index}: {instructions}\n") + parts.append("\n".join(f" {codes[i]}: {item}" for i, item in enumerate(labels))) + parts.append("\n\nReply with one answer code per question, in order.") + text = tokenizer.apply_chat_template( + [{"role": "user", "content": "".join(parts)}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + return _encode(tokenizer, text) + + +def build_seed_canvas( + question_count, + canvas_length, + vocab_size, +): + """Build one deterministic active canvas token per question.""" + if question_count > canvas_length: + raise RequestContractError( + f"{question_count} questions need {question_count} canvas tokens; " + f"DiffusionGemma provides {canvas_length}" + ) + rng = random.Random(DEFAULT_DIFFUSION_SEED) + canvas = [rng.randrange(vocab_size) for _ in range(question_count)] + return canvas, list(range(question_count)) + + +class MlxVlmDiffusionEngine: + """Read all question slots in one DiffusionGemma denoising forward.""" + + def __init__(self, model_id=None, base_url=None, transport=None, compact=False): + self.model_id = resolve_model_id(model_id or "diffusion") + if self.model_id != DIFFUSION_GEMMA_MODEL: + raise ValueError("the mlx-vlm-diffusion backend requires the diffusion model profile") + self.model_path, self.model_revision = resolve_model_snapshot(self.model_id) + self.tokenizer = load_tokenizer(self.model_path) + self.registry_file, registry, self.codes, self.code_token_ids = packaged_code_registry( + self.model_id + ) + if registry.get("model_revision") != self.model_revision: + raise ValueError(f"{self.registry_file.name} does not match the pinned model revision") + if registry.get("tokenizer_sha256") != tokenizer_sha256(self.model_path): + raise ValueError(f"{self.registry_file.name} does not match the pinned tokenizer") + if registry.get("codes") != diffusion_answer_code_entries(self.tokenizer): + raise ValueError(f"{self.registry_file.name} is incomplete or invalid") + + config = json.loads((self.model_path / "config.json").read_text()) + self.canvas_length = int(config["canvas_length"]) + self.vocab_size = int(config["text_config"]["vocab_size"]) + self.base_url = ( + base_url or os.environ.get("SYSTEM_ONE_MLX_VLM_BASE_URL") or DEFAULT_BASE_URL + ).rstrip("/") + self.timeout = float(os.environ.get("SYSTEM_ONE_MLX_VLM_TIMEOUT", DEFAULT_TIMEOUT)) + self.api_key = os.environ.get("SYSTEM_ONE_MLX_VLM_API_KEY") + self.transport = transport or self._post + self.compact = compact + + def output_tokens_for(self, questions): + """A read-only canvas does not commit generated tokens.""" + return 0 + + def evaluate(self, state, questions, template=None): + """Return one answer-code distribution per slot in a shared canvas.""" + if template is not None: + raise RequestContractError("custom templates are not supported by DiffusionGemma") + self._validate_questions(questions) + if self.compact and len(questions) != 1: + raise RequestContractError( + "the compact DiffusionGemma backend requires exactly one question" + ) + input_ids = build_prompt( + self.tokenizer, + state, + questions, + self.codes, + compact=self.compact, + ) + if len(input_ids) > MAX_PROMPT_TOKENS: + raise RequestContractError(f"prompt exceeds the {MAX_PROMPT_TOKENS} token limit") + if len(input_ids) > MAX_TOTAL_INPUT_TOKENS: + raise RequestContractError( + f"request exceeds the {MAX_TOTAL_INPUT_TOKENS} input token limit" + ) + seed_canvas, positions = build_seed_canvas( + len(questions), + self.canvas_length, + self.vocab_size, + ) + slots = [ + { + "position": position, + "token_ids": list(self.code_token_ids[: len(labels)]), + } + for position, (_, labels) in zip(positions, questions) + ] + payload = { + "model": str(self.model_path), + "input_ids": input_ids, + "seed_canvas": seed_canvas, + "slots": slots, + "candidate_only": True, + } + started = time.perf_counter() + response = self.transport("/v1/diffusion/reads", payload) + elapsed_ms = (time.perf_counter() - started) * 1000 + probabilities = self._parse_response( + response, + slots, + len(input_ids), + str(self.model_path), + ) + return probabilities, len(input_ids), elapsed_ms + + def _validate_questions(self, questions): + if not questions: + raise RequestContractError("at least one question is required") + for instructions, labels in questions: + if not labels: + raise RequestContractError(f"{instructions!r} has no options") + if len(labels) > len(self.codes): + raise TooManyOptions( + f"{len(labels)} options for {instructions!r}; " + f"the engine has {len(self.codes)} single-token answer codes" + ) + + def _parse_response( + self, + response, + slots, + input_tokens, + expected_model, + ): + try: + reads = response["reads"] + usage = response["usage"] + reported_tokens = usage["prompt_tokens"] + denoising_steps = usage["denoising_steps"] + candidate_only = usage["candidate_only"] + model = response["model"] + except (KeyError, TypeError) as error: + raise MlxVlmError("MLX-VLM returned an invalid diffusion read response") from error + if model != expected_model: + raise MlxVlmError(f"MLX-VLM served {model!r}, expected pinned model {expected_model!r}") + if denoising_steps != 1: + raise MlxVlmError(f"MLX-VLM performed {denoising_steps!r} denoising steps, expected 1") + if candidate_only is not True: + raise MlxVlmError("MLX-VLM did not confirm candidate-only scoring") + if reported_tokens != input_tokens: + raise MlxVlmError( + "MLX-VLM token count differs from the pinned tokenizer: " + f"expected {input_tokens}, got {reported_tokens}" + ) + if len(reads) != len(slots): + raise MlxVlmError("MLX-VLM returned incomplete diffusion read slots") + + probabilities = [] + for read, slot in zip(reads, slots): + try: + token_ids = read["token_ids"] + logprobs = read["logprobs"] + position = read["position"] + except (KeyError, TypeError) as error: + raise MlxVlmError("MLX-VLM returned an invalid diffusion read slot") from error + if position != slot["position"] or token_ids != slot["token_ids"]: + raise MlxVlmError("MLX-VLM returned a mismatched diffusion read slot") + if len(logprobs) != len(token_ids) or not all( + isinstance(value, (int, float)) and math.isfinite(value) for value in logprobs + ): + raise MlxVlmError("MLX-VLM returned invalid token log probabilities") + probabilities.append(restricted_softmax(logprobs)) + return probabilities + + def _post(self, path, payload): + headers = {"Content-Type": "application/json"} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + request = Request( + f"{self.base_url}{path}", + data=json.dumps(payload).encode("utf-8"), + headers=headers, + method="POST", + ) + try: + with urlopen(request, timeout=self.timeout) as response: + return json.load(response) + except HTTPError as error: + detail = error.read(4096).decode("utf-8", errors="replace") + raise MlxVlmError(f"MLX-VLM returned HTTP {error.code}: {detail}") from error + except (URLError, TimeoutError) as error: + raise MlxVlmError(f"could not reach MLX-VLM at {self.base_url}") from error + except (json.JSONDecodeError, UnicodeDecodeError) as error: + raise MlxVlmError("MLX-VLM returned invalid JSON") from error diff --git a/server/src/system_one_lite/vllm_metal.py b/server/src/system_one_lite/vllm_metal.py new file mode 100644 index 0000000..5941dcf --- /dev/null +++ b/server/src/system_one_lite/vllm_metal.py @@ -0,0 +1,250 @@ +"""System One backend for a separate vllm-metal OpenAI server.""" + +import json +import math +import os +import time +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +from mlx_lm.utils import load_tokenizer + +from .engine import ( + MAX_PROMPT_TOKENS, + MAX_TOTAL_INPUT_TOKENS, + TEMPERATURE, + RequestContractError, + TooManyOptions, + common_prefix_len, + configured_model, + load_code_registry, + packaged_code_registry, + resolve_model_snapshot, +) +from .prompts import chat_filled + +DEFAULT_BASE_URL = "http://127.0.0.1:8000" +DEFAULT_TIMEOUT = 120.0 +MAX_LOGPROB_TOKEN_IDS = 128 + + +class VllmMetalError(RuntimeError): + """The vllm-metal server could not complete a structured read.""" + + +def restricted_softmax(logprobs, temperature=TEMPERATURE): + """Renormalize raw vocabulary log-probabilities over allowed codes.""" + scaled = [value / temperature for value in logprobs] + peak = max(scaled) + weights = [math.exp(value - peak) for value in scaled] + total = sum(weights) + return [value / total for value in weights] + + +class VllmMetalEngine: + """Read answer-code logits through vllm-metal's completions API. + + vllm-metal does not yet implement DiffusionGemma or seeded canvas reads. + This backend sends System One's exact prompt token IDs in one batch. Each + prompt stays independent. Requested answer codes are split into exact + groups of 128 because that is vLLM's per-read limit. + """ + + def __init__(self, model_id=None, base_url=None, transport=None, tokenizer=None): + self.model_id = configured_model(model_id) + if tokenizer is None: + model_path, model_revision = resolve_model_snapshot(self.model_id, tokenizer_only=True) + self.tokenizer = load_tokenizer(model_path) + self.registry_file, self.codes, self.code_token_ids = load_code_registry( + self.model_id, + self.tokenizer, + model_path, + model_revision, + ) + else: + self.tokenizer = tokenizer + self.registry_file, _, self.codes, self.code_token_ids = packaged_code_registry( + self.model_id + ) + self.base_url = ( + base_url or os.environ.get("SYSTEM_ONE_VLLM_BASE_URL") or DEFAULT_BASE_URL + ).rstrip("/") + self.timeout = float(os.environ.get("SYSTEM_ONE_VLLM_TIMEOUT", DEFAULT_TIMEOUT)) + self.api_key = os.environ.get("SYSTEM_ONE_VLLM_API_KEY") + self.transport = transport or self._post + + def output_tokens_for(self, questions): + """Count the internal one-token probes needed for this request.""" + return sum(math.ceil(len(labels) / MAX_LOGPROB_TOKEN_IDS) for _, labels in questions) + + def evaluate(self, state, questions, template=None): + """Return one exact code distribution per independent question.""" + prompts, input_tokens = self._prepare(state, questions, template) + started = time.perf_counter() + self._verify_server_tokenizer() + probabilities, reported_input_tokens = self._read_batches(prompts, questions) + elapsed_ms = (time.perf_counter() - started) * 1000 + if reported_input_tokens != input_tokens: + raise VllmMetalError( + "vllm-metal token count differs from the pinned tokenizer: " + f"expected {input_tokens}, got {reported_input_tokens}" + ) + return probabilities, input_tokens, elapsed_ms + + def _prepare(self, state, questions, template): + if template is not None: + raise RequestContractError("custom templates are not supported by vllm-metal") + if not questions: + raise RequestContractError("at least one question is required") + + prompts = [] + input_tokens = 0 + for instructions, labels in questions: + if not labels: + raise RequestContractError(f"{instructions!r} has no options") + if len(labels) > len(self.codes): + raise TooManyOptions( + f"{len(labels)} options for {instructions!r}; " + f"the engine has {len(self.codes)} single-token answer codes" + ) + prompt = self._prompt_token_ids(state, instructions, labels) + if len(prompt) > MAX_PROMPT_TOKENS: + raise RequestContractError( + f"prompt for {instructions!r} exceeds the {MAX_PROMPT_TOKENS} token limit" + ) + input_tokens += len(prompt) * math.ceil(len(labels) / MAX_LOGPROB_TOKEN_IDS) + if input_tokens > MAX_TOTAL_INPUT_TOKENS: + raise RequestContractError( + f"request exceeds the {MAX_TOTAL_INPUT_TOKENS} input token limit" + ) + prompts.append(prompt) + return prompts, input_tokens + + def _prompt_token_ids(self, state, instructions, labels): + text, marks = chat_filled( + self.tokenizer, + state, + [(instructions, labels)], + codes=self.codes, + ) + prefix = text[: marks[0]] + base = self.tokenizer.encode(prefix) + candidate = self.tokenizer.encode(prefix + self.codes[0]) + cut = common_prefix_len(candidate, base) + full = self.tokenizer.encode(text) + if len(candidate) - cut != 1 or candidate[: cut + 1] != full[: cut + 1]: + raise RequestContractError(f"slot for {instructions!r} is not a single token") + return full[:cut] + + def _verify_server_tokenizer(self): + token_ids = list(self.code_token_ids) + response = self.transport( + "/detokenize", + {"model": self.model_id, "tokens": token_ids}, + ) + try: + remote_text = response["prompt"] + except (KeyError, TypeError) as error: + raise VllmMetalError("vllm-metal returned an invalid detokenize response") from error + if remote_text != self.tokenizer.decode(token_ids): + raise VllmMetalError("vllm-metal does not match the pinned tokenizer") + + def _read_batches(self, prompts, questions): + max_options = max(len(labels) for _, labels in questions) + scores = [{} for _ in questions] + prompt_tokens = 0 + for start in range(0, max_options, MAX_LOGPROB_TOKEN_IDS): + end = min(start + MAX_LOGPROB_TOKEN_IDS, max_options) + active = [ + (index, prompt) + for index, (prompt, (_, labels)) in enumerate(zip(prompts, questions)) + if len(labels) > start + ] + token_ids = list(self.code_token_ids[start:end]) + payload = { + "model": self.model_id, + "prompt": [prompt for _, prompt in active], + "max_tokens": 1, + "temperature": 0, + "logprobs": 0, + "logprob_token_ids": token_ids, + "return_tokens_as_token_ids": True, + } + response = self.transport("/v1/completions", payload) + try: + choices = sorted(response["choices"], key=lambda choice: choice["index"]) + response_model = response["model"] + batch_prompt_tokens = response["usage"]["prompt_tokens"] + except (KeyError, TypeError, ValueError) as error: + raise VllmMetalError("vllm-metal returned an invalid logprobs response") from error + if response_model != self.model_id: + raise VllmMetalError( + f"vllm-metal served {response_model!r}, expected {self.model_id!r}" + ) + if ( + isinstance(batch_prompt_tokens, bool) + or not isinstance(batch_prompt_tokens, int) + or batch_prompt_tokens < 0 + ): + raise VllmMetalError("vllm-metal returned an invalid prompt token count") + prompt_tokens += batch_prompt_tokens + if [choice["index"] for choice in choices] != list(range(len(active))): + raise VllmMetalError("vllm-metal returned incomplete batch choices") + for choice, (question_index, _) in zip(choices, active): + option_count = len(questions[question_index][1]) + requested_ids = list(self.code_token_ids[start : min(end, option_count)]) + scores[question_index].update(self._choice_logprobs(choice, requested_ids)) + + probabilities = [ + restricted_softmax([score[token_id] for token_id in self.code_token_ids[: len(labels)]]) + for score, (_, labels) in zip(scores, questions) + ] + return probabilities, prompt_tokens + + def _choice_logprobs(self, choice, token_ids): + try: + top_logprobs = choice["logprobs"]["top_logprobs"][0] + if not isinstance(top_logprobs, dict): + raise TypeError + by_token_id = { + int(token.removeprefix("token_id:")): logprob + for token, logprob in top_logprobs.items() + if token.startswith("token_id:") + } + except (AttributeError, IndexError, KeyError, TypeError, ValueError) as error: + raise VllmMetalError("vllm-metal returned invalid token logprobs") from error + missing = [token_id for token_id in token_ids if token_id not in by_token_id] + if missing: + raise VllmMetalError( + "vllm-metal omitted requested answer-code token IDs: " + + ", ".join(map(str, missing)) + ) + if any( + isinstance(by_token_id[token_id], bool) + or not isinstance(by_token_id[token_id], (int, float)) + or not math.isfinite(by_token_id[token_id]) + for token_id in token_ids + ): + raise VllmMetalError("vllm-metal returned invalid token logprobs") + return {token_id: by_token_id[token_id] for token_id in token_ids} + + def _post(self, path, payload): + headers = {"Content-Type": "application/json"} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + request = Request( + f"{self.base_url}{path}", + data=json.dumps(payload).encode("utf-8"), + headers=headers, + method="POST", + ) + try: + with urlopen(request, timeout=self.timeout) as response: + return json.load(response) + except HTTPError as error: + detail = error.read(4096).decode("utf-8", errors="replace") + raise VllmMetalError(f"vllm-metal returned HTTP {error.code}: {detail}") from error + except (URLError, TimeoutError) as error: + raise VllmMetalError(f"could not reach vllm-metal at {self.base_url}") from error + except (json.JSONDecodeError, UnicodeDecodeError) as error: + raise VllmMetalError("vllm-metal returned invalid JSON") from error diff --git a/server/tests/integration/test_api.py b/server/tests/integration/test_api.py index bfc7f5e..9d2e3ac 100644 --- a/server/tests/integration/test_api.py +++ b/server/tests/integration/test_api.py @@ -7,11 +7,13 @@ from system_one_lite import api as api_module from system_one_lite.api import app, create_app from system_one_lite.engine import RequestContractError +from system_one_lite.mlx_vlm_diffusion import MlxVlmError from system_one_lite.schemas import ( MAX_QUESTIONS, MAX_REQUEST_BYTES, MAX_STATE_CHARS, ) +from system_one_lite.vllm_metal import VllmMetalError client = None test_app = None @@ -49,6 +51,99 @@ def test_configured_engine_uses_selected_profile(monkeypatch): assert selected == ["larger"] +def test_configured_engine_can_select_vllm_metal(monkeypatch): + selected = [] + monkeypatch.setenv("SYSTEM_ONE_BACKEND", "vllm-metal") + monkeypatch.setenv("SYSTEM_ONE_MODEL", "larger") + monkeypatch.setattr(api_module, "VllmMetalEngine", selected.append) + + assert api_module.configured_engine() is None + assert selected == ["larger"] + + +def test_configured_engine_can_select_mlx_vlm_diffusion(monkeypatch): + selected = [] + monkeypatch.setenv("SYSTEM_ONE_BACKEND", "mlx-vlm-diffusion") + monkeypatch.delenv("SYSTEM_ONE_MODEL", raising=False) + monkeypatch.setattr(api_module, "MlxVlmDiffusionEngine", selected.append) + + assert api_module.configured_engine() is None + assert selected == [None] + + +def test_configured_engine_rejects_unknown_backend(monkeypatch): + monkeypatch.setenv("SYSTEM_ONE_BACKEND", "unknown") + + with pytest.raises(ValueError, match="unknown SYSTEM_ONE_BACKEND"): + api_module.configured_engine() + + +def test_vllm_metal_usage_counts_internal_read_tokens(): + class FakeVllmMetalEngine: + model_id = "test-model" + + def evaluate(self, state, questions): + return [[0.75, 0.25] for _ in questions], 20, 1.0 + + def output_tokens_for(self, questions): + return len(questions) + + request = { + "state": "state", + "questions": { + "first": {"type": "noul", "instructions": "first?"}, + "second": {"type": "noul", "instructions": "second?"}, + }, + } + with TestClient(create_app(engine=FakeVllmMetalEngine())) as vllm_client: + response = vllm_client.post("/evaluate", json=request) + + assert response.status_code == 200 + assert response.json()["usage"] == {"input_tokens": 20, "output_tokens": 2} + + +def test_vllm_metal_failure_becomes_bad_gateway(): + class FailedVllmMetalEngine: + model_id = "test-model" + + def evaluate(self, state, questions): + raise VllmMetalError("server unavailable") + + def output_tokens_for(self, questions): + return len(questions) + + request = { + "state": "state", + "questions": {"first": {"type": "noul", "instructions": "first?"}}, + } + with TestClient(create_app(engine=FailedVllmMetalEngine())) as vllm_client: + response = vllm_client.post("/evaluate", json=request) + + assert response.status_code == 502 + assert response.json()["detail"] == "server unavailable" + + +def test_mlx_vlm_failure_becomes_bad_gateway(): + class FailedDiffusionEngine: + model_id = "test-model" + + def evaluate(self, state, questions): + raise MlxVlmError("diffusion server unavailable") + + def output_tokens_for(self, questions): + return 0 + + request = { + "state": "state", + "questions": {"first": {"type": "noul", "instructions": "first?"}}, + } + with TestClient(create_app(engine=FailedDiffusionEngine())) as diffusion_client: + response = diffusion_client.post("/evaluate", json=request) + + assert response.status_code == 502 + assert response.json()["detail"] == "diffusion server unavailable" + + def test_quickstart_mixed_three_questions(): """The quickstart.md request: Choice + Score + Noul, one pass, full shape.""" r = evaluate( diff --git a/server/tests/tools/test_benchmark_diffusion.py b/server/tests/tools/test_benchmark_diffusion.py new file mode 100644 index 0000000..c92764a --- /dev/null +++ b/server/tests/tools/test_benchmark_diffusion.py @@ -0,0 +1,72 @@ +import json +import math + +from tools.benchmark_diffusion import DiffusionClient, Sample, percentile, run_level + + +class FakeClient: + prompt_tokens = 100 + canvas_tokens = 15 + questions = (("first", ["yes", "no"]),) * 3 + + def call(self, nonce): + return Sample(client_ms=float(nonce + 1), round_trip_ms=float(nonce) + 0.5) + + +class FakeTokenizer: + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return list(text.encode()) + + def apply_chat_template(self, messages, **kwargs): + del kwargs + return messages[0]["content"] + + +class FakeEngine: + tokenizer = FakeTokenizer() + codes = ("A", "B") + canvas_length = 256 + vocab_size = 200 + compact = True + + +def test_percentile_uses_nearest_rank(): + assert percentile([5, 1, 4, 2, 3], 0.95) == 5 + + +def test_state_is_fixed_between_requests(): + assert json.loads(DiffusionClient._state()) == { + "ticket": "Everything is down and we have a demo at noon." + } + + +def test_client_metadata_uses_the_selected_profile(): + client = DiffusionClient(FakeEngine(), (("first", ("yes", "no")),)) + + assert client.canvas_tokens == 1 + + +def test_run_level_reports_throughput_and_latency(): + result = run_level(FakeClient(), concurrency=2, request_count=3, request_index_start=0) + + assert result.concurrency == 2 + assert result.requests == 3 + assert result.decisions_per_request == 3 + assert result.prompt_tokens == 100 + assert result.canvas_tokens == 15 + assert math.isclose(result.decisions_per_second, result.requests_per_second * 3) + assert result.client_median_ms == 2.0 + assert result.client_p95_ms == 3.0 + assert result.round_trip_median_ms == 1.5 + assert result.round_trip_p95_ms == 2.5 + + +def test_run_level_uses_the_client_question_count(): + client = FakeClient() + client.questions = (("first", ["yes", "no"]), ("second", ["yes", "no"])) + + result = run_level(client, concurrency=1, request_count=1, request_index_start=0) + + assert result.decisions_per_request == 2 + assert math.isclose(result.decisions_per_second, result.requests_per_second * 2) diff --git a/server/tests/unit/test_mlx_vlm_diffusion.py b/server/tests/unit/test_mlx_vlm_diffusion.py new file mode 100644 index 0000000..8c71a13 --- /dev/null +++ b/server/tests/unit/test_mlx_vlm_diffusion.py @@ -0,0 +1,251 @@ +import math +from pathlib import Path + +import pytest + +from system_one_lite.engine import RequestContractError +from system_one_lite.mlx_vlm_diffusion import ( + MlxVlmDiffusionEngine, + MlxVlmError, + build_prompt, + build_seed_canvas, +) + + +class FakeTokenizer: + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return [ord(char) % 200 for char in text] + + def apply_chat_template(self, messages, **kwargs): + del kwargs + return "CHAT:" + messages[0]["content"] + + +class TextTokenizer: + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return text + + def apply_chat_template(self, messages, **kwargs): + del kwargs + return messages[0]["content"] + + +def bare_engine(transport, canvas_length=256): + engine = MlxVlmDiffusionEngine.__new__(MlxVlmDiffusionEngine) + engine.model_id = "diffusion" + engine.model_path = Path("/pinned/diffusion") + engine.tokenizer = FakeTokenizer() + engine.codes = ("A", "B", "C") + engine.code_token_ids = (10, 11, 12) + engine.canvas_length = canvas_length + engine.vocab_size = 200 + engine.transport = transport + engine.compact = False + return engine + + +def test_full_prompt_asks_for_answer_codes(): + prompt = build_prompt( + TextTokenizer(), + "state", + [("pick", ["yes", "no"])], + ("A", "B"), + ) + + assert "one answer code per question" in prompt + assert '"- label"' not in prompt + + +def test_evaluate_reads_all_questions_from_one_seeded_canvas(): + captured = {} + + def transport(path, payload): + captured.update(path=path, payload=payload) + return { + "model": "/pinned/diffusion", + "reads": [ + { + "position": slot["position"], + "token_id": slot["token_ids"][0], + "token_ids": slot["token_ids"], + "logprobs": [-0.1 * (index + 1) for index in range(len(slot["token_ids"]))], + } + for slot in payload["slots"] + ], + "usage": { + "prompt_tokens": len(payload["input_ids"]), + "denoising_steps": 1, + "candidate_only": payload["candidate_only"], + }, + } + + engine = bare_engine(transport) + questions = [("first", ["yes", "no"]), ("second", ["a", "b", "c"])] + + probabilities, input_tokens, elapsed_ms = engine.evaluate("state", questions) + + assert captured["path"] == "/v1/diffusion/reads" + assert captured["payload"]["model"] == "/pinned/diffusion" + assert len(captured["payload"]["input_ids"]) == input_tokens + assert len(captured["payload"]["seed_canvas"]) == 2 + assert [slot["position"] for slot in captured["payload"]["slots"]] == [0, 1] + assert "encoder_layers" not in captured["payload"] + assert captured["payload"]["candidate_only"] is True + assert [slot["token_ids"] for slot in captured["payload"]["slots"]] == [ + [10, 11], + [10, 11, 12], + ] + assert len(probabilities) == 2 + assert all(math.isclose(sum(row), 1.0) for row in probabilities) + assert elapsed_ms >= 0 + assert engine.output_tokens_for(questions) == 0 + + +def test_seed_canvas_rejects_more_questions_than_fit(): + with pytest.raises(RequestContractError, match="canvas tokens"): + build_seed_canvas(6, canvas_length=5, vocab_size=200) + + +def test_seed_canvas_fits_the_public_64_question_limit(): + first, positions = build_seed_canvas(64, canvas_length=256, vocab_size=200) + second, second_positions = build_seed_canvas(64, canvas_length=256, vocab_size=200) + + assert first == second + assert positions == second_positions + assert len(first) == 64 + assert positions == list(range(64)) + + +def test_compact_read_uses_one_canvas_token_per_question(): + captured = {} + + def transport(path, payload): + captured.update(path=path, payload=payload) + return { + "model": "/pinned/diffusion", + "reads": [ + { + "position": 0, + "token_id": 11, + "token_ids": [10, 11], + "logprobs": [-2.0, -0.1], + } + ], + "usage": { + "prompt_tokens": len(payload["input_ids"]), + "denoising_steps": 1, + "candidate_only": payload["candidate_only"], + }, + } + + engine = bare_engine(transport) + engine.compact = True + probabilities, _, _ = engine.evaluate("delivered order.", [("late?", ["yes", "no"])]) + + assert captured["payload"]["seed_canvas"] == [163] + assert captured["payload"]["slots"] == [{"position": 0, "token_ids": [10, 11]}] + assert "encoder_layers" not in captured["payload"] + assert probabilities[0][1] > probabilities[0][0] + + +def test_compact_read_rejects_more_than_one_question(): + engine = bare_engine(None) + engine.compact = True + + with pytest.raises(RequestContractError, match="exactly one question"): + engine.evaluate( + "state", + [("first", ["yes", "no"]), ("second", ["yes", "no"])], + ) + + +def test_response_must_match_requested_slot(): + engine = bare_engine(None) + with pytest.raises(MlxVlmError, match="mismatched"): + engine._parse_response( + { + "model": "/pinned/diffusion", + "reads": [ + { + "position": 2, + "token_ids": [10, 11], + "logprobs": [-1.0, -2.0], + } + ], + "usage": { + "prompt_tokens": 3, + "denoising_steps": 1, + "candidate_only": True, + }, + }, + [{"position": 1, "token_ids": [10, 11]}], + 3, + "/pinned/diffusion", + ) + + +def test_response_must_confirm_model_and_one_step(): + engine = bare_engine(None) + base = { + "model": "/pinned/diffusion", + "reads": [ + { + "position": 1, + "token_ids": [10, 11], + "logprobs": [-1.0, -2.0], + } + ], + "usage": { + "prompt_tokens": 3, + "denoising_steps": 1, + "candidate_only": True, + }, + } + + wrong_model = {**base, "model": "another-model"} + with pytest.raises(MlxVlmError, match="expected pinned model"): + engine._parse_response( + wrong_model, + [{"position": 1, "token_ids": [10, 11]}], + 3, + "/pinned/diffusion", + ) + + wrong_steps = {**base, "usage": {**base["usage"], "denoising_steps": 2}} + with pytest.raises(MlxVlmError, match="expected 1"): + engine._parse_response( + wrong_steps, + [{"position": 1, "token_ids": [10, 11]}], + 3, + "/pinned/diffusion", + ) + + +def test_full_read_accepts_server_without_encoder_layer_metadata(): + engine = bare_engine(None) + + probabilities = engine._parse_response( + { + "model": "/pinned/diffusion", + "reads": [ + { + "position": 1, + "token_ids": [10, 11], + "logprobs": [-1.0, -2.0], + } + ], + "usage": { + "prompt_tokens": 3, + "denoising_steps": 1, + "candidate_only": True, + }, + }, + [{"position": 1, "token_ids": [10, 11]}], + 3, + "/pinned/diffusion", + ) + + assert len(probabilities) == 1 + assert math.isclose(sum(probabilities[0]), 1.0) diff --git a/server/tests/unit/test_vllm_metal.py b/server/tests/unit/test_vllm_metal.py new file mode 100644 index 0000000..ac2f3e3 --- /dev/null +++ b/server/tests/unit/test_vllm_metal.py @@ -0,0 +1,301 @@ +import math + +import pytest + +from system_one_lite import vllm_metal as vllm_module +from system_one_lite.engine import DEFAULT_MODEL, RequestContractError, TooManyOptions +from system_one_lite.vllm_metal import ( + MAX_LOGPROB_TOKEN_IDS, + VllmMetalEngine, + VllmMetalError, + restricted_softmax, +) + + +class FakeTokenizer: + marker = '{"answer": "' + + def __init__(self, token_count=12): + self.token_count = token_count + + def apply_chat_template(self, messages, **kwargs): + return "" + messages[0]["content"] + "" + + def encode(self, text): + before, marker, suffix = text.rpartition(self.marker) + if not marker: + return list(range(self.token_count)) + tokens = list(range(self.token_count)) + if suffix: + tokens.append(32) + tokens.extend(ord(character) for character in suffix[1:]) + return tokens + + def decode(self, token_ids): + return ",".join(map(str, token_ids)) + + +def response_for(payload, scores=None, prompt_tokens=None): + token_ids = payload["logprob_token_ids"] + scores = scores or [-float(index) for index in range(len(token_ids))] + return { + "model": DEFAULT_MODEL, + "choices": [ + { + "index": choice_index, + "logprobs": { + "top_logprobs": [ + { + f"token_id:{token_id}": score + for token_id, score in zip(token_ids, scores) + } + ] + }, + } + for choice_index in range(len(payload["prompt"])) + ], + "usage": { + "prompt_tokens": prompt_tokens + if prompt_tokens is not None + else len(payload["prompt"]) * 12 + }, + } + + +def fake_engine(transport, token_count=12): + return VllmMetalEngine( + DEFAULT_MODEL, + transport=transport, + tokenizer=FakeTokenizer(token_count), + ) + + +def test_exact_requested_token_logprobs_are_normalized(): + calls = [] + tokenizer = FakeTokenizer() + + def transport(path, payload): + calls.append((path, payload)) + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + return response_for(payload, [-2.0, -1.0, -3.0]) + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + probabilities, input_tokens, elapsed_ms = engine.evaluate( + "state", [("pick", ["one", "two", "three"])] + ) + + assert probabilities[0] == pytest.approx(restricted_softmax([-2.0, -1.0, -3.0])) + assert input_tokens == 12 + assert elapsed_ms >= 0 + assert [path for path, _ in calls] == ["/detokenize", "/v1/completions"] + payload = calls[-1][1] + assert payload["model"] == DEFAULT_MODEL + assert payload["max_tokens"] == 1 + assert payload["logprob_token_ids"] == list(engine.code_token_ids[:3]) + assert payload["prompt"] == [list(range(12))] + assert engine.output_tokens_for([("pick", ["one", "two", "three"])]) == 1 + + +def test_questions_are_independent_and_token_usage_is_summed(): + payloads = [] + tokenizer = FakeTokenizer(token_count=7) + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + payloads.append(payload) + return response_for(payload, prompt_tokens=14) + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + probabilities, input_tokens, _ = engine.evaluate( + "shared", [("first", ["yes", "no"]), ("second", ["up", "down"])] + ) + + assert len(probabilities) == 2 + assert input_tokens == 14 + assert len(payloads) == 1 + assert len(payloads[0]["prompt"]) == 2 + + +def test_more_than_128_options_are_read_in_exact_chunks(): + tokenizer = FakeTokenizer(token_count=5) + payloads = [] + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + payloads.append(payload) + return response_for(payload, prompt_tokens=5) + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + labels = [f"option {index}" for index in range(MAX_LOGPROB_TOKEN_IDS + 2)] + probabilities, input_tokens, _ = engine.evaluate("state", [("pick", labels)]) + + assert [len(payload["logprob_token_ids"]) for payload in payloads] == [128, 2] + assert len(probabilities[0]) == 130 + assert sum(probabilities[0]) == pytest.approx(1.0) + assert input_tokens == 10 + assert engine.output_tokens_for([("pick", labels)]) == 2 + + +def test_mixed_option_counts_skip_finished_questions_on_later_chunks(): + tokenizer = FakeTokenizer(token_count=5) + batch_sizes = [] + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + batch_sizes.append(len(payload["prompt"])) + return response_for(payload, prompt_tokens=len(payload["prompt"]) * 5) + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + questions = [ + ("small", ["yes", "no"]), + ("large", [f"option {index}" for index in range(129)]), + ] + _, input_tokens, _ = engine.evaluate("state", questions) + + assert batch_sizes == [2, 1] + assert input_tokens == 15 + assert engine.output_tokens_for(questions) == 3 + + +def test_server_tokenizer_mismatch_is_an_error_before_inference(): + calls = [] + + def transport(path, payload): + calls.append(path) + return {"prompt": "wrong"} + + engine = fake_engine(transport) + with pytest.raises(VllmMetalError, match="pinned tokenizer"): + engine.evaluate("state", [("pick", ["yes", "no"])]) + + assert calls == ["/detokenize"] + + +def test_missing_requested_token_is_an_error(): + tokenizer = FakeTokenizer() + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + response = response_for(payload) + response["choices"][0]["logprobs"]["top_logprobs"][0].popitem() + return response + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + with pytest.raises(VllmMetalError, match="omitted requested"): + engine.evaluate("state", [("pick", ["yes", "no"])]) + + +def test_wrong_response_model_is_an_error(): + tokenizer = FakeTokenizer() + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + return {**response_for(payload), "model": "wrong-model"} + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + with pytest.raises(VllmMetalError, match="expected"): + engine.evaluate("state", [("pick", ["yes", "no"])]) + + +@pytest.mark.parametrize("prompt_tokens", ["12", -1, True]) +def test_invalid_prompt_token_count_is_an_error(prompt_tokens): + tokenizer = FakeTokenizer() + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + return response_for(payload, prompt_tokens=prompt_tokens) + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + with pytest.raises(VllmMetalError, match="prompt token count"): + engine.evaluate("state", [("pick", ["yes", "no"])]) + + +@pytest.mark.parametrize("bad_value", [float("nan"), float("inf"), "-1.0", True]) +def test_invalid_requested_logprob_is_an_error(bad_value): + tokenizer = FakeTokenizer() + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + return response_for(payload, scores=[bad_value, -1.0]) + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + with pytest.raises(VllmMetalError, match="invalid token logprobs"): + engine.evaluate("state", [("pick", ["yes", "no"])]) + + +def test_non_mapping_logprobs_are_an_error(): + tokenizer = FakeTokenizer() + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + response = response_for(payload) + response["choices"][0]["logprobs"]["top_logprobs"][0] = [] + return response + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + with pytest.raises(VllmMetalError, match="invalid token logprobs"): + engine.evaluate("state", [("pick", ["yes", "no"])]) + + +def test_contract_checks_happen_before_transport(monkeypatch): + def transport(path, payload): + raise AssertionError("transport should not run") + + engine = fake_engine(transport) + with pytest.raises(RequestContractError, match="at least one"): + engine.evaluate("state", []) + with pytest.raises(RequestContractError, match="custom templates"): + engine.evaluate("state", [("pick", ["yes"])], template=object()) + with pytest.raises(TooManyOptions): + engine.evaluate("state", [("pick", ["x"] * (len(engine.codes) + 1))]) + + monkeypatch.setattr(vllm_module, "MAX_PROMPT_TOKENS", 11) + with pytest.raises(RequestContractError, match="prompt for"): + engine.evaluate("state", [("pick", ["yes", "no"])]) + + +def test_total_token_limit_is_checked_before_transport(monkeypatch): + def transport(path, payload): + raise AssertionError("transport should not run") + + engine = fake_engine(transport, token_count=10) + monkeypatch.setattr(vllm_module, "MAX_TOTAL_INPUT_TOKENS", 19) + + with pytest.raises(RequestContractError, match="request exceeds"): + engine.evaluate( + "state", + [("first", ["yes", "no"]), ("second", ["up", "down"])], + ) + + +def test_prompt_at_input_limit_keeps_one_remote_output_slot(monkeypatch): + tokenizer = FakeTokenizer(token_count=12) + + def transport(path, payload): + if path == "/detokenize": + return {"prompt": tokenizer.decode(payload["tokens"])} + return response_for(payload, prompt_tokens=12) + + engine = VllmMetalEngine(DEFAULT_MODEL, transport=transport, tokenizer=tokenizer) + monkeypatch.setattr(vllm_module, "MAX_PROMPT_TOKENS", 12) + + probabilities, input_tokens, _ = engine.evaluate("state", [("pick", ["yes", "no"])]) + + assert len(probabilities[0]) == 2 + assert input_tokens == 12 + + +def test_restricted_softmax_applies_system_one_temperature(): + probabilities = restricted_softmax([-2.0, -1.0]) + expected_second = 1 / (1 + math.exp(-1 / 0.7)) + + assert probabilities == pytest.approx([1 - expected_second, expected_second]) diff --git a/server/tools/benchmark_diffusion.py b/server/tools/benchmark_diffusion.py new file mode 100644 index 0000000..320abb8 --- /dev/null +++ b/server/tools/benchmark_diffusion.py @@ -0,0 +1,213 @@ +"""Profile warmed System One DiffusionGemma requests through MLX-VLM. + +Start the MLX-VLM server first. Then run this command from ``server/``: + + uv run python -m tools.benchmark_diffusion --json +""" + +import argparse +import json +import math +import platform +import statistics +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import asdict, dataclass + +from system_one_lite.mlx_vlm_diffusion import ( + DEFAULT_BASE_URL, + MlxVlmDiffusionEngine, + build_prompt, + build_seed_canvas, +) + +REQUEST_INDEX_START = 0 +STATE = "Everything is down and we have a demo at noon." +QUESTIONS = ( + ( + "Does the customer need a reply within the hour?", + ("yes", "no"), + ), + ( + "Which team should handle this request?", + ( + "outage — Service outage or availability issue", + "billing — Charges, refunds, or payment methods", + "feature — Feature request or product suggestion", + ), + ), + ( + "What is the customer's tone?", + ("calm", "annoyed", "furious"), + ), +) + + +@dataclass(frozen=True) +class Sample: + client_ms: float + round_trip_ms: float + + +@dataclass(frozen=True) +class Result: + concurrency: int + requests: int + decisions_per_request: int + prompt_tokens: int + canvas_tokens: int + wall_seconds: float + requests_per_second: float + decisions_per_second: float + client_median_ms: float + client_p95_ms: float + round_trip_median_ms: float + round_trip_p95_ms: float + + +def percentile(values, fraction): + ordered = sorted(values) + return ordered[math.ceil(fraction * len(ordered)) - 1] + + +class DiffusionClient: + def __init__(self, engine, questions=QUESTIONS): + self.engine = engine + self.questions = tuple((instructions, list(labels)) for instructions, labels in questions) + state = self._state() + self.prompt_tokens = len( + build_prompt( + engine.tokenizer, + state, + self.questions, + engine.codes, + compact=engine.compact, + ) + ) + canvas, _ = build_seed_canvas( + len(self.questions), + engine.canvas_length, + engine.vocab_size, + ) + self.canvas_tokens = len(canvas) + + @staticmethod + def _state(): + return json.dumps({"ticket": STATE}) + + def call(self, _request_index): + started = time.perf_counter() + probabilities, prompt_tokens, round_trip_ms = self.engine.evaluate( + self._state(), self.questions + ) + client_ms = (time.perf_counter() - started) * 1000 + if prompt_tokens != self.prompt_tokens: + raise RuntimeError("prompt token count changed during the benchmark") + if len(probabilities) != len(self.questions): + raise RuntimeError("engine returned the wrong number of decisions") + if any(abs(sum(row) - 1.0) > 1e-5 for row in probabilities): + raise RuntimeError("engine returned an invalid probability distribution") + return Sample(client_ms=client_ms, round_trip_ms=round_trip_ms) + + +def run_level(client, concurrency, request_count, request_index_start): + started = time.perf_counter() + with ThreadPoolExecutor(max_workers=concurrency) as executor: + samples = list( + executor.map( + client.call, + range(request_index_start, request_index_start + request_count), + ) + ) + wall_seconds = time.perf_counter() - started + request_rate = request_count / wall_seconds + client_times = [sample.client_ms for sample in samples] + round_trip_times = [sample.round_trip_ms for sample in samples] + return Result( + concurrency=concurrency, + requests=request_count, + decisions_per_request=len(client.questions), + prompt_tokens=client.prompt_tokens, + canvas_tokens=client.canvas_tokens, + wall_seconds=wall_seconds, + requests_per_second=request_rate, + decisions_per_second=request_rate * len(client.questions), + client_median_ms=statistics.median(client_times), + client_p95_ms=percentile(client_times, 0.95), + round_trip_median_ms=statistics.median(round_trip_times), + round_trip_p95_ms=percentile(round_trip_times, 0.95), + ) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--base-url", default=DEFAULT_BASE_URL) + parser.add_argument("--timeout", type=float, default=600.0) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--requests", type=int, default=64) + parser.add_argument("--concurrency", default="1,2,4,8,16,32") + parser.add_argument("--questions", type=int, default=len(QUESTIONS)) + parser.add_argument("--profile", choices=("full", "fast"), default="full") + parser.add_argument("--json", action="store_true") + args = parser.parse_args() + if args.warmup < 0 or args.requests < 1: + parser.error("warmup must be non-negative and requests must be positive") + concurrency = [int(value) for value in args.concurrency.split(",")] + if not concurrency or any(value < 1 for value in concurrency): + parser.error("concurrency values must be positive") + if args.questions < 1: + parser.error("questions must be positive") + if args.profile == "fast" and args.questions != 1: + parser.error("the fast profile requires exactly one question") + + engine = MlxVlmDiffusionEngine( + base_url=args.base_url, + compact=args.profile == "fast", + ) + engine.timeout = args.timeout + questions = tuple(QUESTIONS[index % len(QUESTIONS)] for index in range(args.questions)) + client = DiffusionClient(engine, questions) + timed_requests = args.requests * len(concurrency) + for request_index in range(args.warmup): + client.call(REQUEST_INDEX_START + timed_requests + request_index) + + results = [] + request_index = REQUEST_INDEX_START + for level in concurrency: + results.append(run_level(client, level, args.requests, request_index)) + request_index += args.requests + + payload = { + "backend": ("mlx-vlm-diffusion-fast" if args.profile == "fast" else "mlx-vlm-diffusion"), + "model": engine.model_id, + "model_revision": engine.model_revision, + "base_url": engine.base_url, + "host": { + "platform": platform.platform(), + "machine": platform.machine(), + "python": platform.python_version(), + }, + "state": STATE, + "warmup": args.warmup, + "results": [asdict(result) for result in results], + } + if args.json: + print(json.dumps(payload, indent=2)) + return + print(f"model: {engine.model_id}") + print( + f"profile: {args.profile}; decisions/request: {len(client.questions)}; " + f"prompt tokens: {client.prompt_tokens}; canvas tokens: {client.canvas_tokens}" + ) + print("clients req/s decisions/s client p50 client p95 roundtrip p50 roundtrip p95") + for result in results: + print( + f"{result.concurrency:7} {result.requests_per_second:5.2f} " + f"{result.decisions_per_second:11.2f} " + f"{result.client_median_ms:10.1f} {result.client_p95_ms:10.1f} " + f"{result.round_trip_median_ms:13.1f} {result.round_trip_p95_ms:13.1f}" + ) + + +if __name__ == "__main__": + main() diff --git a/server/tools/gen_answer_codes.py b/server/tools/gen_answer_codes.py index dc90541..5b1e9dd 100644 --- a/server/tools/gen_answer_codes.py +++ b/server/tools/gen_answer_codes.py @@ -18,7 +18,9 @@ from system_one_lite.engine import ( DEFAULT_MODEL, + DIFFUSION_GEMMA_MODEL, answer_code_entries, + diffusion_answer_code_entries, resolve_model_id, resolve_model_snapshot, tokenizer_sha256, @@ -49,7 +51,10 @@ def main(): model_id, tokenizer_only=True, revision=args.revision ) tokenizer = load_tokenizer(model_path) - entries = answer_code_entries(tokenizer) + is_diffusion = model_id == DIFFUSION_GEMMA_MODEL + entries = ( + diffusion_answer_code_entries(tokenizer) if is_diffusion else answer_code_entries(tokenizer) + ) output = Path(args.output) if args.output else default_output(model_id) if not output.is_absolute(): output = Path.cwd() / output @@ -58,7 +63,11 @@ def main(): "model": model_id, "model_revision": revision, "tokenizer_sha256": tokenizer_sha256(model_path), - "slot_context": "codes appended after the assistant JSON answer prefix", + "slot_context": ( + "codes occupy one seeded DiffusionGemma canvas token after a list marker" + if is_diffusion + else "codes appended after the assistant JSON answer prefix" + ), "codes": entries, } output.write_text(json.dumps(out, indent=1) + "\n")