import json
import os
import torch
from safetensors import safe_open
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
from PIL import Image
MODEL_PATH = "Gyubeum/Qwen3-VL-8B-Instruct-UI-Genie-scoring"
MAX_LEN = 8192
DEVICE = "cuda"
model = Qwen3VLForConditionalGeneration.from_pretrained(
MODEL_PATH,
torch_dtype=torch.bfloat16,
low_cpu_mem_usage=True,
).to(DEVICE).eval()
processor = AutoProcessor.from_pretrained(MODEL_PATH, max_pixels=1_048_576)
tokenizer = processor.tokenizer
def _load_score_head(model_path: str) -> torch.nn.Linear:
from huggingface_hub import hf_hub_download
index_path = hf_hub_download(model_path, "model.safetensors.index.json")
with open(index_path) as f:
weight_map = json.load(f)["weight_map"]
shard_path = hf_hub_download(model_path, weight_map["score.weight"])
with safe_open(shard_path, framework="pt") as f:
weight = f.get_tensor("score.weight")
head = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False)
with torch.no_grad():
head.weight.copy_(weight)
return head.to(DEVICE, dtype=torch.bfloat16)
score_head = _load_score_head(MODEL_PATH)
def score_action(prefix_messages, response_text, images=None):
"""
Score a GUI agent action.
Args:
prefix_messages: Chat messages up to (not including) the assistant turn.
E.g. [{"role": "system", "content": "..."},
{"role": "user", "content": [{"type":"text","text":"..."},
{"type":"image"}]}]
response_text: The assistant tool-call response to score.
images: List of PIL.Image objects matching image placeholders.
Returns:
float: Scalar BT reward (higher = better action).
"""
structured_msgs = list(prefix_messages) + [
{"role": "assistant", "content": response_text.strip()}
]
full_text = processor.apply_chat_template(
structured_msgs, tokenize=False, add_generation_prompt=False
)
inputs = processor(
text=[full_text],
images=images or None,
truncation=True,
max_length=MAX_LEN,
return_tensors="pt",
).to(DEVICE)
with torch.inference_mode():
outputs = model(**inputs, output_hidden_states=True, return_dict=True)
last_hidden = outputs.hidden_states[-1]
last_idx = inputs["attention_mask"].sum(dim=1) - 1
pooled = last_hidden[0, last_idx[0], :]
reward = score_head(pooled.unsqueeze(0).to(score_head.weight.dtype))
return reward.squeeze().float().item()
screenshot = Image.open("screenshot.png").convert("RGB")
SYSTEM_PROMPT = "You are a helpful assistant."
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": [
{"type": "text", "text": "The user query: tap the search button\n"},
{"type": "image"},
]},
]
action_a = '{"name": "mobile_use", "arguments": {"action": "click", "coordinate": [540, 120]}}'
action_b = '{"name": "mobile_use", "arguments": {"action": "click", "coordinate": [100, 800]}}'
score_a = score_action(messages, action_a, images=[screenshot])
score_b = score_action(messages, action_b, images=[screenshot])
print(f"Action A: {score_a:.4f}")
print(f"Action B: {score_b:.4f}")
print(f"Preferred: {'A' if score_a > score_b else 'B'}")