add serve/system1_vllm_backend.py

serve/system1_vllm_backend.py

@@ -0,0 +1,181 @@
+"""System One server on top of a STOCK vLLM OpenAI server (no custom vLLM code).
+
+Run: python system1_vllm_backend.py --vllm-url http://HOST:8000 --served-model s1 --port 8020
+
+Readout (zero generated content): the prompt is rendered and tokenised locally (byte-identical to training), then ONE
+`/v1/completions` request is sent with prompt=<token ids>, max_tokens=1, allowed_token_ids=<candidate ids>, logprobs=K.
+With the server flag `--logprobs-mode processed_logprobs` the returned logprobs are taken AFTER the allowed-token mask, so
+they are exactly log_softmax over the candidates of the next-token logits at the assistant prefix. The backend's scores
+are those logprobs; the Service's softmax (calibrated_probabilities, T=1) is shift-invariant, so probabilities equal the
+MLX backend's softmax over raw candidate logits. Default raw_logprobs mode would return the top-K of the FULL vocabulary
+(candidates can be missing), so a missing candidate or a non-normalised row raises instead of returning wrong scores.
+Sampling params are pinned (temperature 1, top_p 1, top_k -1) because vLLM otherwise applies the model's generation_config.
+
+Fallback `--readout prompt_logprobs`: one request per candidate with prompt = prefix+[candidate id], prompt_logprobs=0; the
+logprob of the last prompt token is log_softmax(full vocab)[candidate]. Exact too (constant shift), needs no server flag, but
+K prefills (vLLM does not reuse the prefix cache for prompt_logprobs requests) and large payloads; use it as a GPU control."""
+import argparse
+import base64
+import json
+import mimetypes
+import math
+import os
+import urllib.error
+import urllib.request
+
+from system1_client import MAX_CONTEXT_TOKENS, OverlengthError
+from system1_server import make_server
+
+MODEL_ID, MODEL_REVISION = "Qwen/Qwen3.5-9B", "c202236235762e1c871ad0ccb60c8ee5ba337b9a"
+
+
+def candidate_ids(tokenizer, prompt, prefix, keys):
+    """Same rules as MlxBackend.score: each key is exactly one token appended at the assistant prefix; ids distinct."""
+    ids = {}
+    for key in keys:
+        extended = tokenizer.encode(prompt + key, add_special_tokens=False)
+        if extended[:-1] != prefix or tokenizer.decode(extended[-1:]) != key:
+            raise ValueError(f"candidate {key}: not one token at assistant prefix")
+        ids[key] = int(extended[-1])
+    if len(set(ids.values())) != len(keys):
+        raise ValueError("candidate token IDs collide")
+    return ids
+
+
+def load_tokenizer(model=MODEL_ID, revision=MODEL_REVISION):
+    from huggingface_hub import snapshot_download
+    from transformers import AutoTokenizer
+    path = model if os.path.isdir(model) else snapshot_download(model, revision=revision, local_files_only=True)
+    return AutoTokenizer.from_pretrained(path, local_files_only=True, trust_remote_code=False)
+
+
+class VllmBackend:
+    def __init__(self, vllm_url, served_model, tokenizer=None, readout="allowed_logprobs", timeout=60, api_key=None,
+                 image_root=None):
+        if readout not in ("allowed_logprobs", "prompt_logprobs"):
+            raise ValueError(f"unknown readout {readout!r}")
+        self.url = vllm_url.rstrip("/") + "/v1/completions"
+        self.chat_url = vllm_url.rstrip("/") + "/v1/chat/completions"
+        self.image_root = os.path.realpath(image_root) if image_root else None  # request-supplied file paths must live here
+        self.model, self.readout, self.timeout, self.api_key = served_model, readout, timeout, api_key
+        self.tokenizer = tokenizer or load_tokenizer()
+
+    def _post(self, body, url=None):
+        headers = {"Content-Type": "application/json"}
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+        req = urllib.request.Request(url or self.url, json.dumps(body).encode(), headers)
+        try:
+            with urllib.request.urlopen(req, timeout=self.timeout) as r:
+                return json.loads(r.read())["choices"][0]
+        except urllib.error.HTTPError as e:
+            raise RuntimeError(f"vLLM HTTP {e.code}: {e.read()[:500].decode(errors='replace')}")
+
+    @staticmethod
+    def _lp(entry):
+        return float(entry["logprob"] if isinstance(entry, dict) else entry)
+
+    def _data_url(self, ref, trusted):
+        if ref.startswith("data:image/"):
+            return ref
+        path = os.path.realpath(ref)
+        if not trusted and not (self.image_root and os.path.commonpath([path, self.image_root]) == self.image_root):
+            raise ValueError("image file paths in requests need --image-root; send a data URL instead")
+        with open(path, "rb") as f:
+            return f"data:{mimetypes.guess_type(path)[0] or 'image/jpeg'};base64," + base64.b64encode(f.read()).decode()
+
+    @staticmethod
+    def _request_image(messages):
+        """A System One request carries state.image={"type":"image_url","image_url":{"url":..}}. Returns (url, messages with
+        the image replaced by "attached"), or (None, messages) for every text-only request (those never get re-serialised)."""
+        content = messages[-1]["content"]
+        if '"image_url"' not in content:
+            return None, messages
+        payload = json.loads(content)
+        img = payload.get("state", {}).get("image") if isinstance(payload.get("state"), dict) else None
+        if not (isinstance(img, dict) and img.get("type") == "image_url" and isinstance(img.get("image_url", {}).get("url"), str)):
+            return None, messages
+        payload["state"] = {**payload["state"], "image": "attached"}
+        text = json.dumps(payload, ensure_ascii=True, separators=(",", ":")).replace("<", "\\u003c")
+        return img["image_url"]["url"], messages[:-1] + [{**messages[-1], "content": text}]
+
+    def score_image(self, messages, keys, image):
+        """Same readout as score() but via /v1/chat/completions with the image prepended to the user turn.
+        `image` is a trusted local path or a data URL. Candidate ids come from the text-only render (same assistant prefix)."""
+        return self._score_image(messages, keys, self._data_url(image, trusted=True))
+
+    def _score_image(self, messages, keys, data_url):
+        prompt = self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
+        prefix = self.tokenizer.encode(prompt, add_special_tokens=False)
+        if len(prefix) + 1 > MAX_CONTEXT_TOKENS:  # ponytail: text only; image tokens are bounded by the server's own limits
+            raise OverlengthError("prompt exceeds context budget")
+        ids = candidate_ids(self.tokenizer, prompt, prefix, keys)
+        chat = messages[:-1] + [{"role": "user", "content": [
+            {"type": "image_url", "image_url": {"url": data_url}}, {"type": "text", "text": messages[-1]["content"]}]}]
+        choice = self._post({"model": self.model, "messages": chat, "max_tokens": 1, "temperature": 1.0, "top_p": 1.0, "top_k": -1,
+                             "chat_template_kwargs": {"enable_thinking": False}, "allowed_token_ids": list(ids.values()),
+                             "logprobs": True, "top_logprobs": len(ids), "return_tokens_as_token_ids": True}, self.chat_url)
+        top = {e["token"]: e["logprob"] for e in choice["logprobs"]["content"][0]["top_logprobs"]}
+        missing = [k for k, t in ids.items() if f"token_id:{t}" not in top]
+        if missing:
+            raise RuntimeError(f"vLLM returned no logprob for candidates {missing}; start vllm serve with --logprobs-mode processed_logprobs")
+        out = {k: self._lp(top[f"token_id:{t}"]) for k, t in ids.items()}
+        if abs(sum(math.exp(v) for v in out.values()) - 1) > 1e-2:
+            raise RuntimeError("candidate logprobs are not normalised over the candidates; start vllm serve with --logprobs-mode processed_logprobs")
+        return out
+
+    def score(self, messages, keys, image=None):
+        if image is not None:
+            return self.score_image(messages, keys, image)
+        url, messages = self._request_image(messages)
+        if url is not None:
+            return self._score_image(messages, keys, self._data_url(url, trusted=False))
+        prompt = self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
+        prefix = self.tokenizer.encode(prompt, add_special_tokens=False)
+        if len(prefix) + 1 > MAX_CONTEXT_TOKENS:
+            raise OverlengthError("prompt exceeds context budget")
+        ids = candidate_ids(self.tokenizer, prompt, prefix, keys)
+        pin = {"model": self.model, "max_tokens": 1, "temperature": 1.0, "top_p": 1.0, "top_k": -1}
+        if self.readout == "prompt_logprobs":
+            out = {}
+            for key, tid in ids.items():
+                last = self._post({**pin, "prompt": prefix + [tid], "prompt_logprobs": 0})["prompt_logprobs"][-1]
+                if str(tid) not in last:
+                    raise RuntimeError(f"vLLM prompt_logprobs missing candidate {key} (token {tid})")
+                out[key] = self._lp(last[str(tid)])
+            return out
+        choice = self._post({**pin, "prompt": prefix, "allowed_token_ids": list(ids.values()), "logprobs": len(ids),
+                             "return_tokens_as_token_ids": True})
+        top = choice["logprobs"]["top_logprobs"][0]
+        missing = [k for k, t in ids.items() if f"token_id:{t}" not in top]
+        if missing:
+            raise RuntimeError(f"vLLM returned no logprob for candidates {missing}; start vllm serve with --logprobs-mode processed_logprobs")
+        out = {k: self._lp(top[f"token_id:{t}"]) for k, t in ids.items()}
+        if abs(sum(math.exp(v) for v in out.values()) - 1) > 1e-2:
+            raise RuntimeError("candidate logprobs are not normalised over the candidates; start vllm serve with --logprobs-mode processed_logprobs")
+        return out
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--vllm-url", required=True)
+    ap.add_argument("--served-model", required=True, help="vLLM model name or --lora-modules name (e.g. s1)")
+    ap.add_argument("--tokenizer", default=MODEL_ID, help="HF id (local cache only) or directory")
+    ap.add_argument("--tokenizer-revision", default=MODEL_REVISION)
+    ap.add_argument("--readout", choices=("allowed_logprobs", "prompt_logprobs"), default="allowed_logprobs")
+    ap.add_argument("--image-root", help="directory request image file paths may point into (default: data URLs only)")
+    ap.add_argument("--port", type=int, default=8020)
+    ap.add_argument("--host", default="127.0.0.1")
+    ap.add_argument("--model-name", default="system1-9b")
+    a = ap.parse_args()
+    backend = VllmBackend(a.vllm_url, a.served_model, load_tokenizer(a.tokenizer, a.tokenizer_revision), a.readout,
+                          image_root=a.image_root)
+    server = make_server(backend, a.model_name, a.host, a.port)
+    print(f"serving {a.model_name} ({a.readout}) on http://{a.host}:{a.port} -> {a.vllm_url}", flush=True)
+    server.serve_forever()
+
+
+if __name__ == "__main__":
+    main()