| """Hidden-state feature of the hidden-state arms (EXPERIMENT_LOG.md 2026-10-19, tasks 3 and 4a) and its HTTP server. |
| Runs in .train-venv. |
| |
| Feature: Qwen3.5-9B's last-layer hidden state (after the final norm, i.e. the lm_head input; 4096 dims) at the last |
| prompt token of A's grounding question (ctrl_policy.ground(): the frame, GROUND.format(name), thinking off, the answer |
| started as START, rendered by the chat template with continue_final_message as vLLM renders it), prefill only. Batches |
| are right-padded and each row's last real token is gathered. Values travel as base64 of the bf16 bits (encode/decode; |
| hid_policy.decode is the numpy twin), so a logged feature is exactly the one the controller used. |
| |
| Server: one model replica per GPU (threads), a shared queue, each replica takes up to MAX_BATCH waiting requests per |
| forward. POST /feat {"frame": base64 PNG, "name": target name} -> {"feat": ...} and, with --base, also "base": the |
| feature with the adapter disabled (the 4a KL reference input). The port opens once every replica is loaded. |
| Usage: CUDA_VISIBLE_DEVICES=2 .train-venv/bin/python train/hid_server.py --port 9210 [--gpus 0] [--adapter DIR --base] |
| """ |
| import argparse |
| import base64 |
| from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer |
| import io |
| import json |
| from pathlib import Path |
| import queue |
| import sys |
| import threading |
| import time |
|
|
| import torch |
| from PIL import Image |
| from transformers import AutoProcessor |
|
|
| from bt_train import MODULES, load_model |
| from ctrl_policy import START |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.append(str(ROOT / "bench")) |
| from diag_perception_direction import GROUND |
|
|
| MAX_BATCH = 32 |
|
|
|
|
| def processor(): |
| return AutoProcessor.from_pretrained(ROOT / "models/Qwen3.5-9B", local_files_only=True) |
|
|
|
|
| def prompt(proc, name): |
| return proc.apply_chat_template( |
| [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": GROUND.format(name)}]}, |
| {"role": "assistant", "content": START}], |
| tokenize=False, add_generation_prompt=False, continue_final_message=True, enable_thinking=False) |
|
|
|
|
| def batch_inputs(proc, images, names): |
| return dict(proc(text=[prompt(proc, n) for n in names], images=images, padding=True, padding_side="right", |
| return_tensors="pt")) |
|
|
|
|
| def features(model, inputs, device): |
| """(B, 4096) bf16 last-layer hidden states at each row's last prompt token; model: the Qwen3.5 model or its |
| PeftModel (grad follows the caller's mode).""" |
| body = (model.get_base_model() if hasattr(model, "get_base_model") else model).model |
| x = {k: v.to(device) for k, v in inputs.items()} |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| h = body(**x, use_cache=False).last_hidden_state |
| return h[torch.arange(len(h), device=device), x["attention_mask"].sum(1) - 1] |
|
|
|
|
| def encode(f): |
| """bf16 vector -> base64 of its bits.""" |
| return base64.b64encode(f.to(torch.bfloat16).view(torch.int16).cpu().numpy().tobytes()).decode() |
|
|
|
|
| def decode(s): |
| return torch.frombuffer(bytearray(base64.b64decode(s)), dtype=torch.int16).view(torch.bfloat16) |
|
|
|
|
| def png(b64): |
| return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB") |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--port", type=int, required=True) |
| ap.add_argument("--gpus", default="0", help="local cuda indices, one replica each") |
| ap.add_argument("--adapter", type=Path) |
| ap.add_argument("--base", action="store_true", help="also return the feature with the adapter disabled") |
| args = ap.parse_args() |
| proc = processor() |
| jobs = queue.Queue() |
|
|
| def replica(model, device): |
| while True: |
| items = [jobs.get()] |
| while len(items) < MAX_BATCH: |
| try: |
| items.append(jobs.get_nowait()) |
| except queue.Empty: |
| break |
| try: |
| inputs = batch_inputs(proc, [png(i["frame"]) for i in items], [i["name"] for i in items]) |
| with torch.no_grad(): |
| out = {"feat": features(model, inputs, device)} |
| if args.base: |
| with model.disable_adapter(): |
| out["base"] = features(model, inputs, device) |
| for k, item in enumerate(items): |
| item["out"] = {key: encode(f[k]) for key, f in out.items()} |
| except Exception as e: |
| for item in items: |
| item["out"] = {"error": f"{type(e).__name__}: {e}"} |
| for item in items: |
| item["done"].set() |
|
|
| gpus = [int(g) for g in args.gpus.split(",")] |
| for g in gpus: |
| model = load_model(g) |
| if args.adapter: |
| from peft import PeftModel |
| model = PeftModel.from_pretrained(model, args.adapter) |
| model.eval() |
| threading.Thread(target=replica, args=(model, torch.device("cuda", g)), daemon=True).start() |
|
|
| class Handler(BaseHTTPRequestHandler): |
| def log_message(self, *a): |
| pass |
|
|
| def reply(self, code, body): |
| data = json.dumps(body).encode() |
| self.send_response(code) |
| self.send_header("Content-Type", "application/json") |
| self.send_header("Content-Length", str(len(data))) |
| self.end_headers() |
| self.wfile.write(data) |
|
|
| def do_GET(self): |
| self.reply(200, {"gpus": gpus}) |
|
|
| def do_POST(self): |
| item = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) |
| item["done"] = threading.Event() |
| jobs.put(item) |
| item["done"].wait() |
| self.reply(500 if "error" in item["out"] else 200, item["out"]) |
|
|
| server = ThreadingHTTPServer(("0.0.0.0", args.port), Handler) |
| server.daemon_threads = True |
| print(f"{time.strftime('%Y-%m-%d %H:%M:%S')} feature server on port {args.port}, gpus {gpus}, adapter " |
| f"{args.adapter}, base {args.base}", flush=True) |
| server.serve_forever() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
