import jsonimport osimport torchfrom safetensors import safe_openfrom transformers import AutoProcessor, Qwen3VLForConditionalGenerationfrom PIL import Image MODEL_PATH = "Gyubeum/Qwen3-VL-8B-Instruct-UI-Genie-scoring"MAX_LEN = 8192DEVICE = "cuda" # HF tolerates the extra `score.weight` tensor in the checkpoint (warning,# harmless) — it just isn't wired into the base CausalLM class, so we load# it ourselves into a separate linear head below.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()} ] # Mirrors training: chat template closes the assistant turn with # <|im_end|>; the score head reads the hidden state at that position. 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] # (1, seq_len, hidden_dim) last_idx = inputs["attention_mask"].sum(dim=1) - 1 pooled = last_hidden[0, last_idx[0], :] # (hidden_dim,) reward = score_head(pooled.unsqueeze(0).to(score_head.weight.dtype)) return reward.squeeze().float().item() # Example: compare two candidate actionsscreenshot = Image.open("screenshot.png").convert("RGB") SYSTEM_PROMPT = "You are a helpful assistant." # use full mobile_use tool spec in practice 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'}")