google

/

gemma-4-26B-A4B-it

like 1.38k

Follow

Google 64.2k

#!/usr/bin/env python3
"""Reproduce: a prompt ending in a bare multimodal token destabilises gemma-4.

The chat template emits multimodal placeholders from `tool` messages *after*
`format_tool_response_block()` has closed the block with `<tool_response|>`.
When the `tool` message is the last one, the prompt ends with a bare `<|image|>`.

This script pins a tool call with `tool_choice`, so the only correct output is
that tool call with no prose. It reports how often the model does something else.

    vllm serve google/gemma-4-26B-A4B-it \
        --enable-auto-tool-choice --tool-call-parser gemma4 --reasoning-parser gemma4 --trust-request-chat-template

    python repro_gemma4_tool_image.py --base http://localhost:8000/v1 \
        --model google/gemma-4-26B-A4B-it -n 100

    # and with a patched template:
    python repro_gemma4_tool_image.py ... --chat-template patched.jinja

Requires Pillow.
"""
import argparse
import base64
import io
import json
import urllib.error
import urllib.request
from collections import Counter

from PIL import Image, ImageDraw


def make_image_data_url(size=896, cells=8, invert=False):
    """A plain checkerboard.

    Any image reproduces the problem, but the *rate* depends on the image:
    a dark-background board fails ~82% here, a light-background one ~40%.
    Use --invert to switch.
    """
    bg, fg = ((40, 40, 40), (235, 235, 235)) if invert else ((235, 235, 235), (40, 40, 40))
    img = Image.new("RGB", (size, size), bg)
    draw = ImageDraw.Draw(img)
    step = size // cells
    for row in range(cells):
        for col in range(cells):
            if (row + col) % 2:
                draw.rectangle(
                    [col * step, row * step, (col + 1) * step - 1, (row + 1) * step - 1],
                    fill=fg)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()


TOOL = {
    "type": "function",
    "function": {
        "name": "pick_category",
        "description": "Pick the category that best matches the request.",
        "parameters": {
            "type": "object",
            "properties": {
                "reason": {"type": "string", "description": "Why this category."},
                "category": {"type": "string",
                             "enum": ["food", "place", "product", "plant", "other"]},
            },
            "required": ["reason", "category"],
        },
    },
}


def build_messages(data_url, with_image):
    """[system, user(image), assistant(tool_call), tool(text [+ image])]"""
    tool_content = [{"type": "text", "text":
                     "# search results\n\n[1] first result\n[2] second result\n[3] third result"}]
    if with_image:
        tool_content.append({"type": "image_url", "image_url": {"url": data_url}})
    return [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": [{"type": "image_url", "image_url": {"url": data_url}}]},
        {"role": "assistant", "tool_calls": [{
            "id": "call_1", "type": "function",
            "function": {"name": "search", "arguments": json.dumps({"q": "x"})}}]},
        {"role": "tool", "tool_call_id": "call_1", "name": "search", "content": tool_content},
    ]


def stream_once(url, body, api_key):
    """SSE request; returns (content, tool_call_seen) or raises."""
    req = urllib.request.Request(
        url, data=json.dumps({**body, "stream": True}).encode(),
        headers={"Content-Type": "application/json", "Authorization": "Bearer " + api_key})
    content, called = [], False
    with urllib.request.urlopen(req, timeout=300) as r:
        for raw in r:
            line = raw.decode("utf-8", "replace").strip()
            if not line.startswith("data:"):
                continue
            payload = line[5:].strip()
            if payload == "[DONE]":
                break
            try:
                chunk = json.loads(payload)
            except json.JSONDecodeError:
                continue
            for ch in chunk.get("choices") or []:
                delta = ch.get("delta") or {}
                if delta.get("content"):
                    content.append(delta["content"])
                for tc in delta.get("tool_calls") or []:
                    if (tc.get("function") or {}).get("name") == "pick_category":
                        called = True
    return "".join(content).strip(), called


def classify(text, called):
    """Correct = the pinned tool call, and nothing else."""
    if called:
        return "ok" if not text else "tool_call + text"
    return "empty" if not text else "text, no tool call"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", default="http://localhost:8000/v1")
    ap.add_argument("--model", required=True)
    ap.add_argument("--api-key", default="EMPTY")
    ap.add_argument("-n", type=int, default=100)
    ap.add_argument("--temperature", type=float, default=1.0)
    ap.add_argument("--chat-template", help="path to a .jinja file (needs --trust-request-chat-template)")
    ap.add_argument("--no-tool-image", action="store_true", help="control: omit the tool-message image")
    ap.add_argument("--invert", action="store_true", help="dark-background checkerboard (higher failure rate)")
    args = ap.parse_args()

    data_url = make_image_data_url(invert=args.invert)
    body = {
        "model": args.model,
        "messages": build_messages(data_url, not args.no_tool_image),
        "tools": [TOOL],
        "tool_choice": {"type": "function", "function": {"name": "pick_category"}},
        "temperature": args.temperature,
        "max_completion_tokens": 512,
        "stream_options": {"include_usage": True},
    }
    if args.chat_template:
        with open(args.chat_template, encoding="utf-8") as f:
            body["chat_template"] = f.read()

    url = args.base.rstrip("/") + "/chat/completions"
    counts, samples = Counter(), []
    for i in range(args.n):
        try:
            text, called = stream_once(url, body, args.api_key)
            kind = classify(text, called)
        except urllib.error.HTTPError as e:
            kind, text = f"HTTP {e.code}", e.read().decode()[:160]
        counts[kind] += 1
        if kind != "ok" and len(samples) < 5:
            samples.append((kind, text[:160]))
        print(f"\r  {i + 1}/{args.n}  ok={counts['ok']}", end="", flush=True)

    n = sum(counts.values())
    bad = n - counts["ok"]
    print(f"\n\n  tool image: {'omitted' if args.no_tool_image else ('dark' if args.invert else 'light')}"
          f"   template: {args.chat_template or 'server default'}")
    print(f"  incorrect: {bad}/{n}  ({100 * bad / n:.1f}%)")
    for k, c in counts.most_common():
        if k != "ok":
            print(f"    {k:20s} {c:4d}  ({100 * c / n:.1f}%)")
    for kind, text in samples:
        print(f"    [{kind}] {text!r}")


main()