posterreward_lite_inference.py

import json
import os
import argparse
from typing import List

def main():
    parser = argparse.ArgumentParser(description="PosterReward 推理脚本")
    parser.add_argument("--model", type=str, required=True)
    parser.add_argument("--input", type=str)
    parser.add_argument("--prompt", type=str)
    parser.add_argument("--image_path", type=str)
    parser.add_argument("--output_dir", type=str, default="./infer_results")
    parser.add_argument("--gpu", type=str, default="0")
    args = parser.parse_args()

    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    from swift.llm import PtEngine, InferRequest

    print(f"正在加载 Reward 模型: {args.model} ...")
    engine = PtEngine(args.model, max_batch_size=64, task_type='seq_cls', num_labels=1)

    infer_requests = []
    
    def create_request(p, img):
        messages = [
            {"role": "user", "content": f"<image>{p}"},
            {"role": "assistant", "content": ""}
        ]
        return InferRequest(messages=messages, images=[img])

    if args.prompt and args.image_path:
        infer_requests.append(create_request(args.prompt, args.image_path))
    elif args.input:
        with open(args.input, 'r', encoding='utf-8') as f:
            for line in f:
                data = json.loads(line)
                infer_requests.append(create_request(data.get('prompt', ''), data.get('image', '')))

    print(f"开始推理，样本总数: {len(infer_requests)}")
    resp_list = engine.infer(infer_requests)

    print("\n" + "="*30 + " 推理结果 " + "="*30)
    for req, resp in zip(infer_requests, resp_list):
        score = resp.choices[0].message.content
        print(f"Image: {os.path.basename(req.images[0])} | Reward: {score}")

if __name__ == "__main__":
    main()