diff --git a/.gitignore b/.gitignore index 3857937..99ba095 100644 --- a/.gitignore +++ b/.gitignore @@ -23,4 +23,6 @@ target # Local Python worker environment .venv/ __pycache__/ +# Model weights downloaded by the recipes +weights/ .DS_Store diff --git a/README.md b/README.md index 76944aa..05d9d08 100644 --- a/README.md +++ b/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. | diff --git a/recipe/README.md b/recipe/README.md index 4d0bf6c..48d22d5 100644 --- a/recipe/README.md +++ b/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/`. diff --git a/recipe/cua_s1/requirements-text.txt b/recipe/cua_s1/requirements-text.txt new file mode 100644 index 0000000..136ef2b --- /dev/null +++ b/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 diff --git a/recipe/cua_s1/text.md b/recipe/cua_s1/text.md new file mode 100644 index 0000000..1fcfcb1 --- /dev/null +++ b/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 +``` diff --git a/src/frontend/cua_s1_text.py b/src/frontend/cua_s1_text.py new file mode 100644 index 0000000..ef3a330 --- /dev/null +++ b/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 --adapter /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() diff --git a/src/models/cua_s1/README.md b/src/models/cua_s1/README.md index 6a07260..aa0e4be 100644 --- a/src/models/cua_s1/README.md +++ b/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@:`, in An error rejects the whole request. Its body is `{"detail": ""}`, 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: diff --git a/src/models/cua_s1/text/contract.py b/src/models/cua_s1/text/contract.py new file mode 100644 index 0000000..326e758 --- /dev/null +++ b/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, + } diff --git a/src/models/cua_s1/text/model.py b/src/models/cua_s1/text/model.py new file mode 100644 index 0000000..ea0415a --- /dev/null +++ b/src/models/cua_s1/text/model.py @@ -0,0 +1,45 @@ +"""Qwen3.5-4B with the Cua-S1 `text` adapter, loaded and scored as upstream +`cua_s1.four_b.FourBModel` does: an unmerged PEFT adapter, the chat template with its +generation prompt, full logits, and a float32 softmax over the option letters at the +last position. That keeps the probabilities bitwise identical to the reference. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import torch +from peft import PeftModel +from transformers import AutoModelForCausalLM, AutoTokenizer + +from .contract import LETTERS, Question, build_messages + + +class TextModel: + def __init__(self, base: str, adapter: str, device: str, dtype: str) -> None: + # PEFT only warns about keys it cannot place: refuse the multimodal adapter. + config = json.loads((Path(adapter) / "adapter_config.json").read_text()) + if "linear_fc1" in config["target_modules"]: + raise ValueError( + f"{adapter} is the multimodal adapter; pass its text/ directory" + ) + self.tokenizer = AutoTokenizer.from_pretrained(base) + model = AutoModelForCausalLM.from_pretrained( + base, dtype=getattr(torch, dtype), device_map=device + ) + self.model = PeftModel.from_pretrained(model, adapter).eval() + self.letter_ids = self.tokenizer.convert_tokens_to_ids(list(LETTERS)) + + def encode(self, state: str, question: Question): + text = self.tokenizer.apply_chat_template( + build_messages(state, question), tokenize=False, add_generation_prompt=True + ) + return self.tokenizer(text, return_tensors="pt") + + @torch.no_grad() + def score(self, inputs, n_options: int) -> list[float]: + logits = self.model(**inputs.to(self.model.device)).logits[0, -1] + return torch.softmax( + logits[self.letter_ids[:n_options]].float(), dim=-1 + ).tolist() diff --git a/tests/cua_s1/test_text_contract.py b/tests/cua_s1/test_text_contract.py new file mode 100644 index 0000000..870c50f --- /dev/null +++ b/tests/cua_s1/test_text_contract.py @@ -0,0 +1,183 @@ +"""Contract tests without weights or torch. The tokenizer test also runs when +CUA_S1_BASE points to a local Qwen/Qwen3.5-4B directory (tokenizer files only). + +PYTHONPATH=src python -m pytest tests/cua_s1 +""" + +import json +import math +import os + +import pytest + +from models.cua_s1.text.contract import ( + RequestError, + answer, + build_messages, + map_request, + parse_body, +) + + +def body(state="Screen", **question): + q = { + "type": "choice", + "instructions": "Pick one.", + "criteria": {"a": "A", "b": "B"}, + } + q.update(question) + return {"model": "cua-s1-4b-0.2", "state": state, "questions": {"q": q}} + + +def mapped(request): + return map_request(parse_body(json.dumps(request).encode())) + + +def reject(request, status=422): + raw = request if isinstance(request, bytes) else json.dumps(request).encode() + with pytest.raises(RequestError) as info: + map_request(parse_body(raw)) + assert info.value.status == status + return str(info.value) + + +# upstream's libs/cua-driver/examples/jev-use/fixtures/jev-choice-request-v1.json, +# with the chooser's rendered region as `state`, and the user message upstream builds for it +FIXTURE = body( + 'Visual-region-derived observation for capture "capture-fixture-1":\n' + "\"submit-text\": text 'Submit' at (300,240,100,40) confidence=0.96 interactive=true", + instructions="Submit the verified form.", + criteria={ + "submit-form": "Submit using the unique validated visual region.", + "reobserve": "Discard this decision set and obtain a fresh observation.", + "abstain": "Stop without acting if no supplied action is safe.", + }, +) +FIXTURE_USER = ( + "Goal: Submit the verified form.\n\n" + "App: Cua Driver\nTask family: closed-candidate decision\n\n" + "Accessibility tree:\n" + 'Visual-region-derived observation for capture "capture-fixture-1":\n' + "\"submit-text\": text 'Submit' at (300,240,100,40) confidence=0.96 interactive=true\n\n" + "Options:\n" + 'A. Decision "Submit using the unique validated visual region." -> select\n' + 'B. Decision "Discard this decision set and obtain a fresh observation." -> select\n' + 'C. Decision "Stop without acting if no supplied action is safe." -> select\n\n' + "Answer with a single letter." +) + + +def test_fixture_prompt_matches_upstream(): + state, (question,) = mapped(FIXTURE) + system, user = build_messages(state, question) + assert user == {"role": "user", "content": FIXTURE_USER} + assert system["content"].startswith( + "You are a one-pass computer-use decision model." + ) + + +@pytest.mark.skipif(not os.environ.get("CUA_S1_BASE"), reason="set CUA_S1_BASE to run") +def test_tokenizer(): + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(os.environ["CUA_S1_BASE"]) + assert tokenizer.convert_tokens_to_ids(list("ABCDEFGHIJKLMNOPQRSTUVWXYZ")) == list( + range(32, 58) + ) + state, (question,) = mapped(FIXTURE) + text = tokenizer.apply_chat_template( + build_messages(state, question), tokenize=False, add_generation_prompt=True + ) + assert text.endswith("<|im_start|>assistant\n\n") + ids = tokenizer(text)["input_ids"] + assert len(ids) == 218 + assert ids == tokenizer(text, add_special_tokens=False)["input_ids"] + + +def test_goal_left_out_when_empty_or_null(): + for goal in ["", None]: + state, (question,) = mapped(body(instructions=goal)) + assert build_messages(state, question)[1]["content"].startswith( + "App: Cua Driver\n" + ) + + +def test_structured_values_escaping_and_null_label(): + tree = { + "app": "Settings", + "elements": [{"id": "e1", "label": "Location", "on": True}], + } + state, (question,) = mapped( + body( + tree, + instructions={"question": "Which one?"}, + criteria={ + "e1": {"action": "click", "element": "e1"}, + "e2": ["click", "e2"], + "quote": 'Click "Submit"\n(tab\there) C:\\Users', + "save": "点击「保存」", + "abstain": None, + }, + ) + ) + assert state == json.dumps(tree, ensure_ascii=False) + assert question.goal == '{"question": "Which one?"}' + assert question.labels == ( + '{\\"action\\": \\"click\\", \\"element\\": \\"e1\\"}', + '[\\"click\\", \\"e2\\"]', + 'Click \\"Submit\\"\\n(tab\\there) C:\\\\Users', + "点击「保存」", + "abstain", + ) + + +def test_request_errors(): + assert "not supported" in reject(body(type="score", criteria=["low", "high"])) + assert "not supported" in reject(body(type="noul")) + assert "unknown type" in reject(body(type="rank")) + assert "1 to 26 options" in reject(body(criteria={})) + assert "1 to 26 options" in reject(body(criteria={f"o{i}": "x" for i in range(27)})) + assert ( + len(mapped(body(criteria={f"o{i}": "x" for i in range(26)}))[1][0].keys) == 26 + ) + no_instructions = body() + del no_instructions["questions"]["q"]["instructions"] + assert "'instructions' is required" in reject(no_instructions) + for value in [1, 2.5, True]: + reject(body(criteria={"a": value})) + for state in ["", {}, [], None, 3, True]: + reject(body(state)) + assert "'model'" in reject({**body(), "model": "english"}) + many = body() + many["questions"] = {f"q{i}": many["questions"]["q"] for i in range(65)} + reject(many, status=413) + + +@pytest.mark.parametrize( + "raw", + [ + b'{"model": "cua-s1-4b-0.2", "state": {"x": 1, "x": 2}}', + b'{"model": "cua-s1-4b-0.2", "state": NaN}', + b'{"model": "cua-s1-4b-0.2", "state": {"x": 1e400}}', + b'{"model": "cua-s1-4b-0.2", "state": {"x": ' + b"9" * 5000 + b"}}", + b'{"model": "cua-s1-4b-0.2", "state": "\\ud800"}', + b"[" * 100000 + b"]" * 100000, + b"\xff\xfe", + b"\xef\xbb\xbf{}", + b"not json", + b"[1, 2]", + ], +) +def test_malformed_bodies_are_400(raw): + reject(raw, status=400) + + +def test_answer(): + _, (question,) = mapped(body()) + tie = answer(question, [0.5, 0.5]) + assert tie["choice"] == "a" and tie["confidence"] == pytest.approx(0.0, abs=1e-12) + result = answer(question, [0.12, 0.88]) + assert result["choice"] == "b" + assert list(result["probabilities"]) == ["a", "b"] + h = -(0.12 * math.log(0.12) + 0.88 * math.log(0.88)) + assert result["confidence"] == pytest.approx(1 - h / math.log(2)) diff --git a/tests/cua_s1/test_text_server.py b/tests/cua_s1/test_text_server.py new file mode 100644 index 0000000..bc4be9c --- /dev/null +++ b/tests/cua_s1/test_text_server.py @@ -0,0 +1,85 @@ +"""HTTP tests for the worker with a fake model: no weights, no torch.""" + +import json + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") +from fastapi.testclient import TestClient # noqa: E402 + +from frontend.cua_s1_text import build_app # noqa: E402 + + +class Ids: + def __init__(self, n): + self.shape = (1, n) + + +class FakeModel: + def __init__(self, tokens=100, error=None, nan=False): + self.tokens, self.error, self.nan, self.calls = tokens, error, nan, 0 + + def encode(self, state, question): + return {"input_ids": Ids(self.tokens)} + + def score(self, inputs, n_options): + self.calls += 1 + if self.error: + raise self.error + p = [0.1] * (n_options - 1) + [1 - 0.1 * (n_options - 1)] + return [float("nan")] + p[1:] if self.nan else p + + +BODY = { + "model": "cua-s1-4b-0.2", + "state": "Screen", + "questions": { + "_sa": { + "type": "choice", + "instructions": "Pick.", + "criteria": {"_x": "A", "b": "B"}, + } + }, +} + + +def post(model=None, **kwargs): + return TestClient(build_app(model or FakeModel())).post("/v1/systemone", **kwargs) + + +def test_health_and_answer(): + app = build_app(FakeModel()) + assert TestClient(app).get("/health").json()["status"] == "ready" + response = post(json=BODY) + assert response.status_code == 200, response.text + reply = response.json() + assert reply["answers"]["_sa"]["choice"] == "b" + assert list(reply["answers"]["_sa"]["probabilities"]) == ["_x", "b"] + assert reply["usage"] == {"input_tokens": 100, "output_tokens": 0} + + +def test_errors(): + bad = json.loads(json.dumps(BODY)) + bad["questions"]["_sa"]["type"] = "noul" + assert post(json=bad).status_code == 422 + assert post(content=b"{").status_code == 400 + raw = b" " * (4 << 20) + json.dumps(BODY).encode() + assert post(content=iter([raw[:10], raw[10:]])).status_code == 413 + model = FakeModel(tokens=20000) + assert post(model, json=BODY).status_code == 413 and model.calls == 0 + + +@pytest.mark.parametrize( + "model", [FakeModel(error=RuntimeError("CUDA out of memory")), FakeModel(nan=True)] +) +def test_model_failure_is_500(model): + response = post(model, json=BODY) + assert response.status_code == 500 + assert response.json() == {"detail": "inference failed"} + + +def test_warmup(): + model = FakeModel() + build_app(model).state.warmup() + assert model.calls == 1