419 lines
15 KiB
Python
419 lines
15 KiB
Python
import json
|
||
import os
|
||
import re
|
||
import logging
|
||
import concurrent.futures
|
||
import argparse
|
||
from openai import OpenAI
|
||
from dotenv import load_dotenv
|
||
import datetime
|
||
import base64
|
||
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
|
||
# from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
|
||
from qwen_vl_utils import process_vision_info
|
||
|
||
|
||
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,3"
|
||
|
||
# 设置日志
|
||
# 创建logger
|
||
log = logging.getLogger(__name__)
|
||
log.setLevel(logging.INFO)
|
||
|
||
# 创建文件处理器
|
||
file_handler = logging.FileHandler('temp_analysis/exam_run_vision.log')
|
||
file_handler.setLevel(logging.INFO)
|
||
|
||
# 创建控制台处理器
|
||
console_handler = logging.StreamHandler()
|
||
console_handler.setLevel(logging.INFO)
|
||
|
||
# 创建格式器并添加到处理器
|
||
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
||
file_handler.setFormatter(formatter)
|
||
console_handler.setFormatter(formatter)
|
||
|
||
# 将处理器添加到logger
|
||
log.addHandler(file_handler)
|
||
log.addHandler(console_handler)
|
||
|
||
# 加载环境变量
|
||
load_dotenv()
|
||
|
||
PARALLEL_WORKERS=1
|
||
|
||
MODEL_NAME=os.getenv("MODEL_NAME")
|
||
print(f"MODEL_NAME: {MODEL_NAME}")
|
||
|
||
# 配置OpenAI客户端
|
||
# client = OpenAI()
|
||
client = OpenAI(
|
||
api_key=os.getenv("OPENAI_API_KEY"),
|
||
base_url=os.getenv("OPENAI_API_BASE_URL") # 如果使用其他兼容服务，可以设置基础URL
|
||
)
|
||
|
||
# modified version of the task proposer agent prompt from https://arxiv.org/pdf/2502.11357
|
||
|
||
# System prompt
|
||
cot_system_prompt = """
|
||
What does this webpage show? Imagine you are a real user on this webpage. Given the webpage
|
||
screenshot or ocr result and parsed HTML/accessibility tree and the task description, please provide
|
||
the first action towards completing that task.
|
||
|
||
Do the following step by step:
|
||
1. Given the webpage screenshot or ocr result and parsed HTML/accessibility tree, generate the first action
|
||
towards completing that task (in natural language form).
|
||
2. Given the webpage screenshot or ocr result, parsed HTML/accessibility tree, and the natural language
|
||
action, generate the grounded version of that action.
|
||
|
||
ACTION SPACE: Your action space is: ['click [element ID]', 'type [element ID] [content]',
|
||
'select [element ID] [content of option to select]', 'scroll [up]', 'scroll [down]', and 'stop'].
|
||
Action output should follow the syntax as given below:
|
||
click [element text]: This action clicks on an element with a specific text on the webpage.
|
||
type [element ID] [content]: Use this to type the content into the field with id. By default, the
|
||
"Enter" key is pressed after typing. Both the content and the ID should be within square braces
|
||
as per the syntax.
|
||
select [element ID] [content of option to select]: Select an option from a dropdown menu. The
|
||
content of the option to select should be within square braces. When you get (select an option)
|
||
tags from the accessibility tree, you need to select the serial number (element_id) corresponding
|
||
to the select tag, not the option, and select the most likely content corresponding to the option as
|
||
input.
|
||
scroll [down]: Scroll the page down.
|
||
scroll [up]: Scroll the page up.
|
||
|
||
IMPORTANT:
|
||
To be successful, it is important to STRICTLY follow the below rules:
|
||
Action generation rules:
|
||
1. You should generate a single atomic action at each step.
|
||
2. The action should be an atomic action from the given vocabulary - click, type, select, scroll
|
||
(up or down), or stop.
|
||
3. The arguments to each action should be within square braces. For example, "click [127]",
|
||
"type [43] [content to type]", "scroll [up]", "scroll [down]".
|
||
4. The natural language form of action (corresponding to the field "action_in_natural_language")
|
||
should be consistent with the grounded version of the action (corresponding to the field "grounded
|
||
_action"). Do NOT add any additional information in the grounded action. For example, if a
|
||
particular element ID is specified in the grounded action, a description of that element must be
|
||
present in the natural language action.
|
||
5. If the type action is selected, the natural language form of action ("action_in_natural_language") should always specify the actual text to be typed.
|
||
6. You should issue a "stop" action if the current webpage asks to log in or for credit card
|
||
information.
|
||
7. To input text, there is NO need to click the textbox first, directly type content. After typing,
|
||
the system automatically hits the 'ENTER' key.
|
||
8. STRICTLY Avoid repeating the same action (click/type) if the webpage remains unchanged.
|
||
You may have selected the wrong web element.
|
||
9. Do NOT use quotation marks in the action generation.
|
||
|
||
OUTPUT FORMAT:
|
||
Please give a short analysis of the screenshot, parsed
|
||
HTML/accessibility tree, then put your answer within ``` ```, for example,
|
||
"In summary, the proposed task and the corresponding action is: ```{
|
||
"action_in_natural_language": "<ACTION_IN_NATURAL_LANGUAGE>:str",
|
||
"grounded_action": "<ACTION>:str"}```
|
||
"""
|
||
|
||
# User prompt
|
||
cot_user_prompt = """
|
||
[
|
||
{{
|
||
"type": "text",
|
||
"text": "Website URL: {INIT_URL}\nParsed HTML/Accessibility Tree: {A11Y_TREE}\nTask description: {TASK_DESCRIPTION}"
|
||
}},
|
||
{{
|
||
"type": "image_url",
|
||
"image_url": {{
|
||
"url": f"file://{SCREENSHOT}"
|
||
}}
|
||
}}
|
||
]
|
||
"""
|
||
|
||
def load_qwen_model():
|
||
"""
|
||
加载Qwen2VL模型和处理器
|
||
"""
|
||
global qwen_model, processor
|
||
|
||
# 设置本地模型路径
|
||
model_path = "/data1/yuyr/Qwen2.5-VL-7B-Instruct"
|
||
# model_path = "/home/yuyr/.cache/huggingface/hub/models--Qwen--Qwen2.5-VL-3B-Instruct/snapshots/1b989f2c63999d7344135894d3cfa8f494116743/"
|
||
# model_path = "/data1/yuyr/models--bytedance-research--UI-TARS-7B-DPO/snapshots/727b0df39207dafc6cf211a61f29d84b7659c39c/"
|
||
|
||
try:
|
||
log.info("正在加载Qwen2VL模型和处理器...")
|
||
|
||
# 加载模型
|
||
qwen_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
||
# qwen_model = Qwen2VLForConditionalGeneration.from_pretrained(
|
||
model_path, torch_dtype="auto", device_map="cuda:1"
|
||
# model_path, torch_dtype="auto", device_map="auto"
|
||
)
|
||
log.info("Qwen2VL模型加载完成")
|
||
|
||
# 加载processor
|
||
processor = AutoProcessor.from_pretrained(model_path)
|
||
log.info("处理器加载完成")
|
||
|
||
log.info(f"Qwen2VL模型和处理器加载完成")
|
||
|
||
return True
|
||
except Exception as e:
|
||
log.error(f"Qwen2VL模型加载失败: {e}")
|
||
return False
|
||
|
||
|
||
def call_api(messages):
|
||
"""使用openai库调用API接口"""
|
||
try:
|
||
log.info(f"call llm messages ")
|
||
response = client.chat.completions.create(
|
||
messages=messages,
|
||
model=MODEL_NAME
|
||
)
|
||
log.info(f"call llm response: {response}")
|
||
return response
|
||
except Exception as e:
|
||
log.error(f"API调用出错: {e}")
|
||
return None
|
||
|
||
def call_qwen_model(messages):
|
||
"""
|
||
调用Qwen2VL模型
|
||
"""
|
||
# 准备推理输入
|
||
try:
|
||
text = processor.apply_chat_template(
|
||
messages, tokenize=False, add_generation_prompt=True
|
||
)
|
||
image_inputs, video_inputs = process_vision_info(messages)
|
||
inputs = processor(
|
||
text=[text],
|
||
images=image_inputs,
|
||
videos=video_inputs,
|
||
padding=True,
|
||
return_tensors="pt",
|
||
)
|
||
inputs = inputs.to(qwen_model.device)
|
||
|
||
# 推理：生成输出
|
||
log.info("正在生成模型回复...")
|
||
generated_ids = qwen_model.generate(**inputs, max_new_tokens=1024)
|
||
generated_ids_trimmed = [
|
||
out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
|
||
]
|
||
output_text = processor.batch_decode(
|
||
generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
|
||
)[0]
|
||
|
||
log.info(f"模型输出: {output_text}")
|
||
return output_text
|
||
except Exception as e:
|
||
log.error(f"模型推理失败: {e}")
|
||
return None
|
||
|
||
|
||
def extract_action(response_text):
|
||
"""从API响应中提取action_in_natural_language和grounded_action"""
|
||
# 使用正则表达式提取JSON部分
|
||
match = re.search(r'```\s*{\s*(.+?)\s*}\s*```', response_text, re.DOTALL)
|
||
if match:
|
||
try:
|
||
# 构建完整的JSON字符串并解析
|
||
json_str = "{" + match.group(1) + "}"
|
||
action_data = json.loads(json_str)
|
||
return action_data.get("action_in_natural_language"), action_data.get("grounded_action")
|
||
except json.JSONDecodeError:
|
||
log.error(f"无法解析JSON: {match.group(1)}")
|
||
return None, None
|
||
|
||
def check_answer(grounded_action, answer):
|
||
try:
|
||
log.info(f"grounded_action: {grounded_action}, answer: {answer}")
|
||
"""检查grounded_action是否匹配answer"""
|
||
if not grounded_action or not grounded_action.startswith("click"):
|
||
return False
|
||
|
||
# 提取element ID
|
||
match = re.search(r'click \[(\d+)\]', grounded_action)
|
||
if not match:
|
||
return False
|
||
|
||
element_id = match.group(1)
|
||
# 将answer拆分为列表（逗号分隔）
|
||
answer_ids = [id.strip() for id in answer.split(",")]
|
||
|
||
# 检查element ID是否在answer列表中
|
||
return element_id in answer_ids
|
||
except Exception as e:
|
||
log.error(f"检查答案出错: {e}")
|
||
return False
|
||
|
||
def process_item(item):
|
||
"""处理单个测试项目"""
|
||
item_num = item["num"]
|
||
item_id = item["id"]
|
||
page_id = item["page_id"]
|
||
url = item["url"]
|
||
task_description = item["task"]
|
||
answer = item["expected_answer"]
|
||
answer_text = item["answer_text"]
|
||
trajectory_id = item["trajectory_id"]
|
||
trajectory_step_num = item["trajectory_step_num"]
|
||
page_child_num = item["page_child_num"]
|
||
|
||
log.info(f"处理ID: {item_id}, URL: {url}")
|
||
|
||
image_path = f"screenshots/{page_id}_{page_child_num}.png"
|
||
if not os.path.exists(image_path):
|
||
log.error(f"找不到文件: {image_path}")
|
||
return None
|
||
|
||
def encode_image_to_base64(image_path):
|
||
"""
|
||
将图片编码为base64字符串
|
||
"""
|
||
try:
|
||
with open(image_path, "rb") as image_file:
|
||
return base64.b64encode(image_file.read()).decode('utf-8')
|
||
except Exception as e:
|
||
log.error(f"图片编码失败: {e}")
|
||
return None
|
||
|
||
SCREENSHOT = image_path
|
||
|
||
# 构建API请求
|
||
messages = [
|
||
{"role": "system", "content": cot_system_prompt},
|
||
{"role": "user", "content": cot_user_prompt.format(
|
||
INIT_URL=url,
|
||
A11Y_TREE="",
|
||
SCREENSHOT=SCREENSHOT, # 这里没有提供截图OCR结果
|
||
TASK_DESCRIPTION=task_description
|
||
)}
|
||
]
|
||
|
||
log.info(f"task_description: {task_description}")
|
||
log.info(f"answer: {answer}, answer_text: {answer_text}")
|
||
|
||
# 调用API
|
||
# response = call_api(messages)
|
||
response = call_qwen_model(messages)
|
||
|
||
print(f"response: {response}")
|
||
|
||
if response:
|
||
content = response
|
||
# if MODEL_NAME == "qwen/qwq-32b:free":
|
||
# reasoning_content = response.choices[0].message.reasoning
|
||
# else:
|
||
# reasoning_content = response.choices[0].message.reasoning_content
|
||
|
||
# log.info(f"reasoning_content: {reasoning_content}")
|
||
log.info(f"content: {content}")
|
||
|
||
# 提取action
|
||
action_nl, grounded_action = extract_action(content)
|
||
|
||
log.info(f"action_nl: {action_nl}, grounded_action: {grounded_action}")
|
||
|
||
# 检查答案，使用文字检查
|
||
is_correct = check_answer(grounded_action, answer_text)
|
||
log.info(f"is_correct: {is_correct}")
|
||
|
||
# 记录结果
|
||
result = {
|
||
"num": item_num,
|
||
"id": item_id,
|
||
"trajectory_id": trajectory_id,
|
||
"trajectory_step_num": trajectory_step_num,
|
||
"page_id": page_id,
|
||
"url": url,
|
||
"task": task_description,
|
||
"expected_answer": answer,
|
||
"thinking": reasoning_content,
|
||
"raw_content": content,
|
||
"action_nl": action_nl,
|
||
"grounded_action": grounded_action,
|
||
"is_correct": is_correct,
|
||
"model": MODEL_NAME,
|
||
"timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
}
|
||
|
||
log.info(f"ID: {item_id}")
|
||
log.info(f"任务: {task_description}")
|
||
log.info(f"动作: {grounded_action}")
|
||
log.info(f"是否正确: {is_correct}")
|
||
log.info("-" * 50)
|
||
|
||
return result
|
||
else:
|
||
log.error(f"API调用失败: {response}")
|
||
return None
|
||
|
||
def main():
|
||
load_qwen_model()
|
||
|
||
# 运行完整测试
|
||
# 读取exam.json
|
||
with open("temp_analysis/test.json", "r") as f:
|
||
exam_data = json.load(f)
|
||
|
||
results = []
|
||
|
||
total_items = len(exam_data) + len(results)
|
||
completed = len(results)
|
||
success_count = len(results)
|
||
fail_count = 0
|
||
|
||
log.info(f"开始测试，需要执行 {total_items} 个任务, 已经成功 {success_count} 个任务")
|
||
|
||
# 使用线程池并发处理
|
||
with concurrent.futures.ThreadPoolExecutor(max_workers=PARALLEL_WORKERS) as executor:
|
||
# 提交所有任务
|
||
future_to_item = {executor.submit(process_item, item): item for item in exam_data}
|
||
|
||
# 收集结果
|
||
for future in concurrent.futures.as_completed(future_to_item):
|
||
result = future.result()
|
||
|
||
# 使用文件锁防止并发写入冲突
|
||
completed += 1
|
||
# 重新读取最新结果,避免覆盖其他线程的写入
|
||
try:
|
||
with open("temp_analysis/results_vl.json", "r") as f:
|
||
results = json.load(f)
|
||
except FileNotFoundError:
|
||
results = []
|
||
|
||
if result:
|
||
results.append(result)
|
||
if result["is_correct"]:
|
||
success_count += 1
|
||
else:
|
||
fail_count += 1
|
||
else:
|
||
fail_count += 1
|
||
|
||
# 打印当前进度
|
||
progress = (completed / total_items) * 100
|
||
log.info(f"进度: {progress:.2f}% ({completed}/{total_items}) - 成功: {success_count}, 失败: {fail_count}")
|
||
|
||
log.info(f"save results to temp_analysis/results_vl.json")
|
||
# 保存结果
|
||
with open("temp_analysis/results_vl.json", "w") as f:
|
||
json.dump(results, f, indent=2)
|
||
|
||
|
||
# 计算正确率
|
||
accuracy = success_count / total_items if total_items > 0 else 0
|
||
|
||
log.info(f"测试完成! 总计: {total_items}题，正确: {success_count}题，错误: {fail_count}题，正确率: {accuracy:.2%}")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
|
||
|
||
|
||
|
||
|
||
|