103 lines
3.0 KiB
Python
103 lines
3.0 KiB
Python
|
import os
|
|||
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
|
|||
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|||
|
|
|
|||
|
|
from PIL import Image
|
|||
|
|
import requests
|
|||
|
|
import torch
|
|||
|
|
from torchvision import io
|
|||
|
|
from typing import Dict
|
|||
|
|
from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
|
|||
|
|
import pickle
|
|||
|
|
import re
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
from peft import PeftModel
|
|||
|
|
|
|||
|
|
# Load the model in half-precision on the available device(s)
|
|||
|
|
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
|||
|
|
"Qwen2.5-VL-3B", torch_dtype=torch.bfloat16, device_map="auto", attn_implementation="flash_attention_2"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 加载LoRA权重，添加load_in_8bit=False和device_map参数
|
|||
|
|
# model = PeftModel.from_pretrained(
|
|||
|
|
# model,
|
|||
|
|
# "hybrid_train_output/checkpoint-100",
|
|||
|
|
# load_in_8bit=False,
|
|||
|
|
# device_map="auto",
|
|||
|
|
# is_trainable=False,
|
|||
|
|
# assign = True
|
|||
|
|
# )
|
|||
|
|
|
|||
|
|
# 确保模型处于评估模式
|
|||
|
|
model.eval()
|
|||
|
|
|
|||
|
|
processor = AutoProcessor.from_pretrained("Qwen2.5-VL-3B")
|
|||
|
|
|
|||
|
|
# Image
|
|||
|
|
# image = Image.open("Tesla.jpg")
|
|||
|
|
|
|||
|
|
# 定义提示文本
|
|||
|
|
text_prompt = (
|
|||
|
|
"<|im_start|>system\n"
|
|||
|
|
"You are a helpful assistant.<|im_end|>\n"
|
|||
|
|
"<|im_start|>user\n"
|
|||
|
|
"<|vision_start|><|image_pad|><|vision_end|>"
|
|||
|
|
"Please tell me the brand of the product in the picture between labels <answer/> and </answer> "
|
|||
|
|
"and explain the reason between labels <thinking/> and </thinking>"
|
|||
|
|
"<|im_end|>\n"
|
|||
|
|
"<|im_start|>assistant"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 加载测试数据
|
|||
|
|
with open("../work/bal_data/test_data.pkl", "rb") as f:
|
|||
|
|
test_data = pickle.load(f)
|
|||
|
|
|
|||
|
|
# 批处理大小
|
|||
|
|
batch_size = 20
|
|||
|
|
correct = 0
|
|||
|
|
total = 0
|
|||
|
|
|
|||
|
|
# 遍历测试数据
|
|||
|
|
for i in tqdm(range(0, len(test_data), batch_size)):
|
|||
|
|
# 准备当前批次的数据
|
|||
|
|
batch = test_data[i:i+batch_size]
|
|||
|
|
batch_images = [item['image'] for item in batch]
|
|||
|
|
batch_brands = [item['brand'] for item in batch]
|
|||
|
|
batch_prompts = [text_prompt] * len(batch_images)
|
|||
|
|
|
|||
|
|
# 模型处理
|
|||
|
|
inputs = processor(
|
|||
|
|
text=batch_prompts,
|
|||
|
|
images=batch_images,
|
|||
|
|
padding=True,
|
|||
|
|
return_tensors="pt"
|
|||
|
|
)
|
|||
|
|
inputs = inputs.to("cuda")
|
|||
|
|
|
|||
|
|
# 生成输出
|
|||
|
|
output_ids = model.generate(**inputs, max_new_tokens=128)
|
|||
|
|
generated_ids = [
|
|||
|
|
output_ids[len(input_ids):]
|
|||
|
|
for input_ids, output_ids in zip(inputs.input_ids, output_ids)
|
|||
|
|
]
|
|||
|
|
output_texts = processor.batch_decode(
|
|||
|
|
generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 提取预测的品牌名称并比较
|
|||
|
|
for pred_text, true_brand in zip(output_texts, batch_brands):
|
|||
|
|
# 使用正则表达式提取<answer>标签中的内容
|
|||
|
|
match = re.search(r'<answer>(.*?)</answer>', pred_text)
|
|||
|
|
if match:
|
|||
|
|
pred_brand = match.group(1).strip().lower()
|
|||
|
|
true_brand = true_brand.lower()
|
|||
|
|
|
|||
|
|
# 比较预测结果
|
|||
|
|
if pred_brand == true_brand:
|
|||
|
|
correct += 1
|
|||
|
|
total += 1
|
|||
|
|
|
|||
|
|
# 计算并输出准确率
|
|||
|
|
accuracy = correct / total if total > 0 else 0
|
|||
|
|
print(f"准确率: {accuracy:.2%} ({correct}/{total})")
|