cua_s1: add a text worker that loads Qwen3.5-4B through Transformers

.gitignore

@@ -23,4 +23,6 @@ target
 # Local Python worker environment
 .venv/
 __pycache__/
+# Model weights downloaded by the recipes
+weights/
 .DS_Store

README.md

@@ -40,7 +40,7 @@ Implementation code lives under `src/`; recipes and documentation stay at the re
 
 | Directory | Responsibility |
 | --- | --- |
-| [`src/frontend/`](src/frontend/) | Rust serving code and the small engine interface. |
+| [`src/frontend/`](src/frontend/) | Rust serving code, Python worker adapters, and the small engine interface. |
 | [`src/models/`](src/models/) | Model implementations, one directory per model: preprocessing, batching, state, execution, and output processing. |
 | [`src/backends/cuda/`](src/backends/cuda/) | NVIDIA GPU operations and kernel integration. |
 | [`src/backends/metal/`](src/backends/metal/) | Apple GPU operations and kernel integration. |

recipe/README.md

@@ -2,6 +2,8 @@
 
 - [Laya text worker](laya/README.md): start the external Python worker, connect the
   Rust frontend and compare direct and proxied responses.
+- [Cua-S1 4B 0.2 text worker](cua_s1/text.md): download the pinned weights, start
+  the worker and connect the Rust frontend.
 
 Recipes contain setup, launch commands and examples. Reusable implementation code
 belongs under `src/`.

recipe/cua_s1/requirements-text.txt

@@ -0,0 +1,14 @@
+# Versions match upstream's `four-b` lock (trycua/cua libs/cua-s1/python/uv.lock
+# at 0e75660ce4c2edda519e0c795fa3ad98abf4e76f), the reference environment in
+# src/models/cua_s1/README.md.
+torch==2.14.0
+transformers==5.17.0
+tokenizers==0.23.2
+peft==0.21.0
+accelerate==1.15.0
+safetensors==0.8.0
+huggingface-hub==1.32.0
+jinja2==3.1.6
+# HTTP serving.
+fastapi==0.141.1
+uvicorn==0.54.0

recipe/cua_s1/text.md

@@ -0,0 +1,46 @@
+# Cua-S1 4B 0.2 text worker
+
+This recipe runs the Cua-S1 4B 0.2 `text` adapter through Transformers and PEFT behind the Rust frontend. It is the correctness reference for native execution. The model is in [`src/models/cua_s1/text/`](../../src/models/cua_s1/text/), the HTTP worker in [`src/frontend/cua_s1_text.py`](../../src/frontend/cua_s1_text.py), and [`src/models/cua_s1/README.md`](../../src/models/cua_s1/README.md) documents the contract. Only `choice` questions are supported.
+
+Run the commands from the repository root, on Linux with an NVIDIA GPU and Python 3.12. The pinned versions match the upstream reference environment:
+
+```sh
+python3.12 -m venv .venv
+.venv/bin/python -m pip install -r recipe/cua_s1/requirements-text.txt
+.venv/bin/hf download Qwen/Qwen3.5-4B \
+  --revision 851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a --local-dir weights/Qwen3.5-4B
+.venv/bin/hf download cua-ai/cua-s1-4b-0.2 \
+  --revision 16818868b0cc7813808aae4e87b417657046ab79 --local-dir weights/cua-s1-4b-0.2
+```
+
+Upstream's `libs/cua-s1/ci/fetch_pinned_weights.py --dest weights --verify-only` (in [trycua/cua](https://github.com/trycua/cua) at `0e75660ce4c2edda519e0c795fa3ad98abf4e76f`) checks every downloaded file against upstream's lock.
+
+Start the worker, which runs one warmup decision before it listens, then the frontend:
+
+```sh
+PYTHONPATH=src .venv/bin/python -m frontend.cua_s1_text \
+  --base weights/Qwen3.5-4B --adapter weights/cua-s1-4b-0.2/text --port 8000
+cargo build --release --locked
+OMNI_JEV_BIND=127.0.0.1:8080 OMNI_JEV_BACKEND_URL=http://127.0.0.1:8000 ./target/release/omni-jev
+```
+
+Requests run one at a time. Bodies over 4 MiB, more than 64 questions, or a prompt over 16,384 tokens get `413`. The worker computes logits for every position, as upstream does, so memory grows with prompt length: the 15,446-token test input peaked at about 21.3 GiB in bfloat16.
+
+```sh
+curl http://127.0.0.1:8080/v1/systemone \
+  -H 'Content-Type: application/json' \
+  -d '{"model":"cua-s1-4b-0.2","state":"Dialog: Delete 3 files permanently? Buttons: Delete, Cancel","questions":{"pick":{"type":"choice","instructions":"Keep the files.","criteria":{"delete":"Click Delete","cancel":"Click Cancel"}}}}'
+```
+
+On an RTX 6000 Ada in bfloat16, the response is:
+
+```json
+{"model":"cua-ai/cua-s1-4b-0.2@16818868b0cc7813808aae4e87b417657046ab79:text","answers":{"pick":{"type":"choice","choice":"cancel","probabilities":{"delete":0.0024726232513785362,"cancel":0.9975274205207825},"confidence":0.9750249565060322}},"usage":{"input_tokens":153,"output_tokens":0}}
+```
+
+The tests need neither weights nor a GPU; with `CUA_S1_BASE=weights/Qwen3.5-4B` they also check the tokenizer:
+
+```sh
+.venv/bin/python -m pip install pytest httpx
+PYTHONPATH=src .venv/bin/python -m pytest tests/cua_s1
+```

src/frontend/cua_s1_text.py

@@ -0,0 +1,110 @@
+"""HTTP worker for Cua-S1 4B 0.2 (`text` adapter): `GET /health` and `POST /v1/systemone`.
+
+PYTHONPATH=src python -m frontend.cua_s1_text --base <dir> --adapter <dir>/text
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import sys
+import traceback
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+
+# FastAPI reads the handler annotations at runtime, so `Request` must be a
+# module-level name while `from __future__ import annotations` is in effect.
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+
+from models.cua_s1.text.contract import (
+    MODEL_ID,
+    RequestError,
+    answer,
+    map_request,
+    parse_body,
+)
+
+MAX_BODY_BYTES = 4 << 20
+MAX_PROMPT_TOKENS = 16384
+WARMUP = (
+    b'{"model": "cua-s1-4b-0.2", "state": "Dialog: Update installed.", "questions": {"q":'
+    b' {"type": "choice", "instructions": "Close it.", "criteria": {"ok": "OK", "wait": "Wait"}}}}'
+)
+
+
+def build_app(model: Any) -> FastAPI:
+    app = FastAPI()
+    pool = ThreadPoolExecutor(max_workers=1)  # one forward pass at a time
+
+    def decide(raw: bytes) -> dict[str, Any]:
+        state, questions = map_request(parse_body(raw))
+        encoded = [model.encode(state, q) for q in questions]
+        tokens = [int(x["input_ids"].shape[1]) for x in encoded]
+        for q, n in zip(questions, tokens):  # before any forward pass
+            if n > MAX_PROMPT_TOKENS:
+                raise RequestError(
+                    f"question {q.name!r}: {n} prompt tokens, over {MAX_PROMPT_TOKENS}",
+                    413,
+                )
+        answers = {
+            q.name: answer(q, model.score(x, len(q.keys)))
+            for q, x in zip(questions, encoded)
+        }
+        return {
+            "model": MODEL_ID,
+            "answers": answers,
+            "usage": {"input_tokens": sum(tokens), "output_tokens": 0},
+        }
+
+    @app.get("/health")
+    def health():
+        return {"status": "ready", "model": MODEL_ID}
+
+    @app.post("/v1/systemone")
+    async def systemone(request: Request):
+        raw = bytearray()
+        async for chunk in request.stream():
+            raw += chunk
+            if len(raw) > MAX_BODY_BYTES:
+                return JSONResponse({"detail": "request body too large"}, 413)
+        try:
+            # JSONResponse, not FastAPI's encoder, which drops keys starting with "_sa".
+            return JSONResponse(
+                await asyncio.get_running_loop().run_in_executor(pool, decide, raw)
+            )
+        except RequestError as error:
+            return JSONResponse({"detail": str(error)}, error.status)
+        except Exception:
+            traceback.print_exc(file=sys.stderr)
+            return JSONResponse({"detail": "inference failed"}, 500)
+
+    # One decision on the worker thread before listening, so the first request does not
+    # pay for first-call setup there.
+    app.state.warmup = lambda: pool.submit(decide, WARMUP).result()
+    return app
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument("--base", required=True, help="local Qwen/Qwen3.5-4B directory")
+    parser.add_argument(
+        "--adapter", required=True, help="local text/ adapter directory"
+    )
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float32"])
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8000)
+    args = parser.parse_args()
+
+    import uvicorn
+
+    from models.cua_s1.text.model import TextModel
+
+    app = build_app(TextModel(args.base, args.adapter, args.device, args.dtype))
+    app.state.warmup()
+    uvicorn.run(app, host=args.host, port=args.port, log_level="warning")
+
+
+if __name__ == "__main__":
+    main()

src/models/cua_s1/README.md

@@ -2,7 +2,7 @@
 
 This directory owns Cua-S1 4B 0.2 ([#10](https://github.com/ThinkFlowLab/system1-omni/issues/10)): request mapping, prompt construction, adapter selection, execution, and the answer-letter readout. This page records the pinned upstream revisions, the inference contract an implementation must match, and how its outputs will be compared with the upstream reference.
 
-Status: planned; nothing is implemented or validated yet. The first target is the `text` adapter on CUDA, starting with a worker that loads the model directly through Hugging Face Transformers and PEFT. The `multimodal` adapter is deferred; see [Not covered yet](#not-covered-yet).
+Status: a reference worker for the `text` adapter loads the model through Hugging Face Transformers and PEFT: [`text/`](text/), served by [`src/frontend/cua_s1_text.py`](../../frontend/cua_s1_text.py), with setup in [`recipe/cua_s1/text.md`](../../../recipe/cua_s1/text.md). It is the correctness reference for native execution, which starts with a minimal `text` prefill path on CUDA. The `multimodal` adapter is deferred; see [Not covered yet](#not-covered-yet).
 
 ## Pinned revisions
 
@@ -76,19 +76,19 @@ The response `model` is `cua-ai/cua-s1-4b-0.2@<adapter revision>:<modality>`, in
 
 An error rejects the whole request. Its body is `{"detail": "<message>"}`, as the LAYA worker returns, and the message names the problem.
 
-The status is `400` when the body is not a usable JSON object: invalid JSON or UTF-8, `NaN` or `Infinity`, a lone surrogate such as `\ud800`, nesting too deep to parse, or a key repeated in any object.
+The status is `400` when the body is not a usable JSON object: invalid JSON or UTF-8, `NaN`, `Infinity` or a number out of range, a lone surrogate such as `\ud800`, nesting too deep to parse, or a key repeated in any object.
 
 The status is `422` when a well-formed request cannot be answered:
 
 - a `score` or `noul` question, since the adapters were trained only on closed-option choices;
 - a question without an `instructions` field (`null` is allowed), or a `choice` with no options or more than 26 options;
 - a `criteria` value that is a number or a boolean;
-- an empty `state`;
+- an empty `state` (`""`, `{}` or `[]`);
 - a `model` other than `cua-s1-4b-0.2`.
 
 ## Validation
 
-**Inputs.** The fixed input set is upstream's two checked-in fixtures, converted to `/v1/systemone` requests with the chooser's rendered regions as `state`, plus `/v1/systemone` choice requests that will be checked in with the worker. These cover 1 to 26 options, short and long states, string and structured `state`, `instructions` and `criteria`, `null` criteria, non-ASCII text, and text that spells a special token. Each input is scored once per configuration.
+**Inputs.** The fixed input set is upstream's two checked-in fixtures, converted to `/v1/systemone` requests with the chooser's rendered regions as `state`, plus `/v1/systemone` choice requests. These cover 1 to 26 options, short and long states, string and structured `state`, `instructions` and `criteria`, `null` criteria, non-ASCII text, and text that spells a special token. Each input is scored once per configuration.
 
 **Tolerances.** These are declared before any comparison is run:
 

src/models/cua_s1/text/contract.py

@@ -0,0 +1,160 @@
+"""Request mapping, prompts and answers for Cua-S1 4B 0.2, following
+`src/models/cua_s1/README.md`. No torch imports, so it can be tested without weights.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+from dataclasses import dataclass
+from typing import Any
+
+MODEL_NAME = "cua-s1-4b-0.2"
+MODEL_ID = "cua-ai/cua-s1-4b-0.2@16818868b0cc7813808aae4e87b417657046ab79:text"
+LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+MAX_QUESTIONS = 64
+
+# The system message and the user message layout are copied from trycua/cua at
+# 0e75660ce4c2edda519e0c795fa3ad98abf4e76f (`libs/cua-s1/python/src/cua_s1/four_b.py`
+# and `libs/cua-driver/examples/jev-use/python/decision_models.py`).
+#
+# MIT License
+#
+# Copyright (c) 2025 Cua AI, Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+SYSTEM_PROMPT = (
+    "You are a one-pass computer-use decision model. You are shown the "
+    "current state of a screen and a fixed, closed list of candidate "
+    "(element, action) options, each given a single letter. Choose exactly "
+    "one option: the single best next action to take. Answer with ONLY that "
+    "option's letter -- no words, no punctuation, no explanation."
+)
+
+
+class RequestError(ValueError):
+    def __init__(self, message: str, status: int = 422) -> None:
+        super().__init__(message)
+        self.status = status
+
+
+@dataclass(frozen=True)
+class Question:
+    name: str
+    goal: str
+    keys: tuple[str, ...]
+    labels: tuple[str, ...]
+
+
+def _unique_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
+    obj = {}
+    for key, value in pairs:
+        if key in obj:
+            raise RequestError(f"duplicate key {key!r}", 400)
+        obj[key] = value
+    return obj
+
+
+def parse_body(raw: bytes) -> dict[str, Any]:
+    try:
+        body = json.loads(raw.decode(), object_pairs_hook=_unique_keys)
+        # NaN, Infinity, numbers out of range and lone surrogates fail here.
+        json.dumps(body, ensure_ascii=False, allow_nan=False).encode()
+    except (ValueError, RecursionError) as error:
+        if isinstance(error, RequestError):
+            raise
+        raise RequestError(f"request body is not valid JSON: {error}", 400) from error
+    if not isinstance(body, dict):
+        raise RequestError("request body must be a JSON object", 400)
+    return body
+
+
+def _text(value: Any, where: str) -> str:
+    """A string as is; an object or array as Python's json.dumps writes it."""
+    if not isinstance(value, (str, dict, list)):
+        raise RequestError(f"{where} must be a string, an object or an array")
+    return value if isinstance(value, str) else json.dumps(value, ensure_ascii=False)
+
+
+def map_request(body: dict[str, Any]) -> tuple[str, list[Question]]:
+    if body.get("model") != MODEL_NAME:
+        raise RequestError(f"'model' must be {MODEL_NAME!r}")
+    if body.get("state") in ("", {}, []):
+        raise RequestError("'state' must not be empty")
+    state = _text(body.get("state"), "'state'")
+    questions = body.get("questions")
+    if not isinstance(questions, dict) or not questions:
+        raise RequestError("'questions' must be a non-empty object")
+    if len(questions) > MAX_QUESTIONS:
+        raise RequestError(f"more than {MAX_QUESTIONS} questions", 413)
+    mapped = []
+    for name, q in questions.items():
+        where = f"question {name!r}"
+        if not isinstance(q, dict):
+            raise RequestError(f"{where} must be an object")
+        if q.get("type") in ("score", "noul"):
+            raise RequestError(f"{where}: type {q['type']!r} is not supported")
+        if q.get("type") != "choice":
+            raise RequestError(f"{where}: unknown type {q.get('type')!r}")
+        if "instructions" not in q:
+            raise RequestError(f"{where}: 'instructions' is required")
+        goal = "" if q["instructions"] is None else _text(q["instructions"], where)
+        criteria = q.get("criteria")
+        if not isinstance(criteria, dict) or not 1 <= len(criteria) <= len(LETTERS):
+            raise RequestError(
+                f"{where}: 'criteria' must be an object with 1 to 26 options"
+            )
+        labels = tuple(
+            json.dumps(
+                key if value is None else _text(value, f"{where}: {key!r}"),
+                ensure_ascii=False,
+            )[1:-1]
+            for key, value in criteria.items()
+        )
+        mapped.append(Question(name, goal, tuple(criteria), labels))
+    return state, mapped
+
+
+def build_messages(state: str, question: Question) -> list[dict[str, str]]:
+    options = "\n".join(
+        f'{letter}. Decision "{label}" -> select'
+        for letter, label in zip(LETTERS, question.labels)
+    )
+    user = (
+        (f"Goal: {question.goal}\n\n" if question.goal else "")
+        + "App: Cua Driver\nTask family: closed-candidate decision\n\n"
+        + f"Accessibility tree:\n{state}\n\nOptions:\n{options}\n\nAnswer with a single letter."
+    )
+    return [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": user},
+    ]
+
+
+def answer(question: Question, probabilities: list[float]) -> dict[str, Any]:
+    """The Jev choice answer; ties go to the earliest option. `confidence` is the
+    normalized entropy `1 - H(p) / ln(n)`, as the LAYA worker reports it."""
+    n = len(probabilities)
+    entropy = -sum(p * math.log(p) for p in probabilities if p > 0)
+    return {
+        "type": "choice",
+        "choice": question.keys[max(range(n), key=probabilities.__getitem__)],
+        "probabilities": dict(zip(question.keys, probabilities)),
+        "confidence": max(0.0, 1 - entropy / math.log(n)) if n > 1 else 1.0,
+    }