import torchfrom PIL import Imagefrom transformers import AutoProcessor, Qwen3VLForConditionalGeneration MODEL_PATH = "Gyubeum/Qwen3-VL-8B-Instruct-UI-Genie-SFT-128K"# SYSTEM_PROMPT, goal, history and tool_call_response come from the agent step being judged.JUDGE_PROMPT = ( "Was the agent's action above correct given the current screen state? " "Answer with exactly one token: <|+|> for correct, <|-|> for wrong.") model = Qwen3VLForConditionalGeneration.from_pretrained( MODEL_PATH, torch_dtype=torch.bfloat16, device_map="cuda",).eval()processor = AutoProcessor.from_pretrained(MODEL_PATH, max_pixels=1_048_576)tok = processor.tokenizerpos_id = tok.convert_tokens_to_ids("<|+|>")neg_id = tok.convert_tokens_to_ids("<|-|>") messages = [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": [ {"type": "text", "text": f"The user query: {goal}\nTask progress (...): {history}"}, {"type": "image", "image": Image.open("screenshot.png").convert("RGB")}, ]}, {"role": "assistant", "content": tool_call_response}, {"role": "user", "content": JUDGE_PROMPT},]inputs = processor.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt",).to(model.device) with torch.no_grad(): logits = model(**inputs).logits[0, -1]top = logits.argmax().item()hard_score = 1.0 if top == pos_id else 0.0 if top == neg_id else 0.5soft_score = torch.softmax(logits[[pos_id, neg_id]].float(), dim=-1)[0].item()