Transformers推理框架

# 导入transformers库中的图像文本到文本模型和处理器类
from transformers import AutoModelForImageTextToText, AutoProcessor

# # 定义要使用的模型名称，会自动从从huggingface仓库下载
# MODEL_NAME = "Qwen/Qwen3-VL-8B-Instruct"
MODEL_NAME = "/home/ubuntu/ai_model/Qwen3-VL-8B-Instruct" # 使用本地模型目录，避免联网下载

# 默认：在可用设备上加载模型
# 从预训练模型加载图像文本到文本的生成模型
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_NAME,  # 使用指定的模型名称
    dtype="auto",  # 自动选择数据类型（float16/float32等）
    # attn_implementation="flash_attention_2", # 多图或视频场景，启用flash_attention_2以提升速度并节省内存
    device_map="auto",  # 自动将模型分配到可用的GPU/CPU设备
)

# 加载与模型对应的处理器，用于处理输入（图像/视频/文本）和输出（文本）
processor = AutoProcessor.from_pretrained(MODEL_NAME)

# 构建对话消息列表，定义用户输入的内容
messages = [
    {
        "role": "user",  # 角色为用户
        "content": [  # 内容列表，可以包含多种类型（视频、图像、文本）
            {
                "type": "video",  # 内容类型为视频
                "video": "data_video/WPT_1.mp4",  # 视频文件路径
            },
            {
                "type": "text",  # 内容类型为文本
                # 用户提示词：要求模型在视频的适当时间点生成评论，返回JSON格式
                "text": "在适当的时候给出或犀利或幽默的像真人的评论。返回json格式，key有timestamp和comment，timestamp是视频的秒数，comment是评论",
            },
        ],
    }
]

# 推理前准备
# 使用处理器将消息转换为模型可接受的输入格式
inputs = processor.apply_chat_template(
    messages,  # 传入消息列表
    tokenize=True,  # 对文本进行分词处理
    add_generation_prompt=True,  # 添加生成提示，告诉模型开始生成回复
    return_dict=True,  # 返回字典格式
    return_tensors="pt",  # 返回PyTorch张量格式
)
# 将输入数据移动到模型所在的设备（GPU或CPU）
inputs = inputs.to(model.device)

# 推理并生成输出
# 使用模型生成文本，max_new_tokens限制生成的最大token数量
generated_ids = model.generate(**inputs, max_new_tokens=512)
# 去除输入部分的token，只保留新生成的token ID
generated_ids_trimmed = [
    out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
]
# 将token ID解码为可读的文本
output_text = processor.batch_decode(
    generated_ids_trimmed,  # 传入修剪后的token ID列表
    skip_special_tokens=True,  # 跳过特殊token（如<pad>、<eos>等）
    clean_up_tokenization_spaces=False,  # 不清理分词后的空格
)
# 打印模型生成的输出文本
print(output_text)