qwenjev / scripts /serve_system2.py
tchbcb's picture
QwenJev: multimodal-retrofitted NanoJev for ARC-AGI-3 (initial skeleton)
3e04895 verified
Raw
History Blame Contribute Delete
2.81 kB
"""Serve the System-2 backbone (Qwen3.5-9B) as an OpenAI-compatible endpoint.
Preferred path - vLLM (recommended by the Qwen3.5 recipes, single GPU):
vllm serve Qwen/Qwen3.5-9B --port 8000
This script is the fallback when vLLM is unavailable: it wraps transformers
in a minimal /v1/chat/completions server. Both expose the same contract the
System2Controller expects (mode="openai").
"""
from __future__ import annotations
import argparse
import base64
import io
import re
def extract_messages(payload: dict) -> tuple[list[dict], dict]:
messages = payload.get("messages", [])
gen_kwargs = {
"temperature": payload.get("temperature", 0.2),
"max_new_tokens": payload.get("max_tokens", 2048),
}
return messages, gen_kwargs
def decode_images(messages: list[dict]) -> list:
"""Pull base64 PNGs out of OpenAI-style image_url content blocks."""
images = []
for msg in messages:
content = msg.get("content")
if isinstance(content, list):
for part in content:
if part.get("type") == "image_url":
url = part["image_url"]["url"]
b64 = re.sub("^data:image/\\w+;base64,", "", url)
from PIL import Image
images.append(Image.open(io.BytesIO(base64.b64decode(b64))))
return images
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--model", default="Qwen/Qwen3.5-9B")
ap.add_argument("--port", type=int, default=8000)
args = ap.parse_args()
import torch
from flask import Flask, jsonify, request
from transformers import AutoProcessor, AutoModelForImageTextToText
processor = AutoProcessor.from_pretrained(args.model)
model = AutoModelForImageTextToText.from_pretrained(
args.model, torch_dtype=torch.bfloat16, device_map="auto"
)
app = Flask(__name__)
@app.post("/v1/chat/completions")
def completions():
payload = request.get_json(force=True)
messages, gen = extract_messages(payload)
text = processor.apply_chat_template(
messages, add_generation_prompt=True, tokenize=False
)
inputs = processor(text=text, images=decode_images(messages) or None,
return_tensors="pt").to(model.device)
out = model.generate(**inputs, **gen)
completion = processor.batch_decode(out[:, inputs["input_ids"].shape[1]:],
skip_special_tokens=True)[0]
return jsonify({
"id": "qwenjev-s2",
"choices": [{"message": {"role": "assistant", "content": completion},
"finish_reason": "stop"}],
})
app.run(host="127.0.0.1", port=args.port)
if __name__ == "__main__":
main()