跳到主要内容
AI教程

使用 WeMM-Embedding 生成多模态嵌入

学习安装 WeMM-Embedding、选择模型、生成多模态向量、缩减向量维度、部署推理服务器并运行 MMEB-v3 评估流程。

WeMM-Embedding 多模态嵌入生成

什么是 WeMM-Embedding?

WeMM-Embedding 是由微信视觉团队开发的一系列通用多模态嵌入模型。它可以将文本、图像、视频、视觉文档和交错式多模态输入转换为统一的向量表示。这些向量可用于需要跨不同模态表示和比较内容的工作流。

WeMM-Embedding performance overview

每个模型都从专用 <embedding> token 位置的最后一层隐藏状态中获取嵌入,然后进行 L2 归一化。目前不支持音频输入。

主要特性

  • 统一的多模态表示:同一模型系列可处理文本、图像、视频、视觉文档和交错式输入。
  • 多种模型规模:可选择 2B、4B 和 9B 参数版本。
  • Matryoshka Representation Learning:当向量存储或检索效率至关重要时,可使用受支持的较短向量维度。
  • 两种推理接口:同时提供 Transformers 和 Sentence Transformers 的示例。
  • 生产环境服务:代码仓库包含适用于 vLLM 和 SGLang 的配置及封装脚本。
  • 评估工具:内置的 MMEB-v3 流程支持单节点和多节点推理。

选择模型和维度

Hugging Face 上的模型库包含三个版本。请根据可用的计算资源和所需的完整嵌入维度选择模型。

  • WeMM-Embedding-2B:支持 64、128、256、512、1024 和 2048 维。
  • WeMM-Embedding-4B:支持 64、128、256、512、1024 和 2560 维。
  • WeMM-Embedding-9B:支持 64、128、256、512、1024、2048 和 4096 维。

使用受支持的较小维度可以减少向量存储空间和下游计算量。例如,README 显示,在 MMEB-v2 上,2B 模型使用 256 维向量时,可保留完整维度下图像和视频性能的 98.7%。

安装与设置

1. 获取代码仓库和模型

在代码仓库的本地副本中操作,然后从对应的 Hugging Face 页面下载模型;如果相关示例支持,也可以将其配置为使用 Hugging Face 模型标识符。

2. 安装依赖项

安装项目指定版本的依赖项:

pip install -r requirements.txt

3. 固定 Transformers 版本以确保可复现性

项目建议使用 transformers==5.2.0 进行推理和结果复现,因为较新版本的预处理行为可能有所不同。

pip install transformers==5.2.0

运行示例前,请准备好模型、测试图像和测试视频的路径。

使用 Transformers 的基本方法

Transformers 示例会分别为文本、图像和视频输入生成独立的嵌入。使用本地模型路径和受支持的输出维度运行:

python examples/transformers_inference.py \
  --model /path/to/WeMM-Embedding-2B \
  --image /path/to/image.jpg \
  --video /path/to/video.mp4 \
  --dimension 2048

请将占位路径替换为系统中的实际文件路径。对于 2B 模型,2048 是完整向量维度。如果希望使用模型的完整嵌入维度而不显式指定,请省略 --dimension

python examples/transformers_inference.py \
  --model /path/to/WeMM-Embedding-2B \
  --image /path/to/image.jpg \
  --video /path/to/video.mp4

使用 Sentence Transformers 的基本方法

Sentence Transformers 示例通过 SentenceTransformer.encode() 处理文本、图像和视频输入。它还支持通过 --dimension 选择 Matryoshka 维度。

python examples/sentence_transformers_inference.py \
  --model /path/to/WeMM-Embedding-2B \
  --image /path/to/image.jpg \
  --video /path/to/video.mp4 \
  --dimension 2048

由于 SentenceTransformer 可以直接加载模型,因此可以将本地路径替换为 Hugging Face 模型标识符:

python examples/sentence_transformers_inference.py \
  --model tencent/WeMM-Embedding-2B \
  --image /path/to/image.jpg \
  --video /path/to/video.mp4 \
  --dimension 2048

正确使用 Matryoshka 嵌入

Matryoshka 嵌入允许你将完整向量缩短为所选模型支持的某个维度。截取最后一个轴,然后再次对缩短后的向量进行归一化:

embedding = torch.nn.functional.normalize(
    embedding[..., :d],
    dim=-1
)

仅可将 d 设置为模型所列出的维度。第二次归一化非常重要,因为截断会改变向量的范数。使用项目提供的命令行示例时,可通过 --dimension 选择所需的 Matryoshka 维度。

使用 vLLM 提供模型服务

代码仓库显示测试所用的 vLLM 版本为 0.27.0。启动 pooling 服务器,并提供模型的嵌入聊天模板:

MODEL_PATH=/path/to/WeMM-Embedding-2B
vllm serve "$MODEL_PATH" \
  --runner pooling \
  --chat-template "$MODEL_PATH/embedding_chat_template.jinja"

scripts/serve_vllm.sh 中提供了等效的封装脚本。

使用 SGLang 提供模型服务

项目文档注明测试所用的 SGLang 版本为 0.5.9。启动嵌入服务器前,请应用项目提供的视频补丁:

MODEL_PATH=/path/to/WeMM-Embedding-2B
python scripts/patch_sglang_video.py
python -m sglang.launch_server \
  --model-path "$MODEL_PATH" \
  --is-embedding \
  --enable-precise-embedding-interpolation

代码仓库还提供了 scripts/serve_sglang.sh,可作为一键运行的封装脚本。

运行 MMEB-v3 评估

mmeb_v3_eval/ 目录包含用于生成已公布 MMEB-v3 结果的评估代码。它基于官方 VLM2Vec 流程,并增加了 WeMM-Embedding 主干网络、对齐的数据集指令、批量推理、多节点多 GPU 执行和 64 帧视频采样。

首先进入评估目录并下载数据:

cd mmeb_v3_eval
DATA_ROOT=/path/to/MMEB-V3 bash scripts/download_data.sh

然后使用模型、数据集和输出路径运行评估:

MODEL_PATH=/path/to/WeMM-Embedding-2B \
DATA_BASEDIR=/path/to/MMEB-V3 \
OUTPUT_DIR=exps/wemm_embedding \
bash scripts/run_eval.sh

有关文档中说明的数据设置以及单节点或多节点命令,请参阅 mmeb_v3_eval/README.md

进阶提示

  1. 确保维度与模型匹配:不要为 2B 模型请求 2560 维,也不要为 4B 模型请求 4096 维。只能使用所选 checkpoint 列出的维度。
  2. 重新归一化截断后的向量:如果手动缩短完整嵌入,之后务必再次进行 L2 归一化。
  3. 确保预处理可复现:复现结果时,请使用推荐的 Transformers 版本,因为不同版本的预处理行为可能发生变化。
  4. 使用合适的加载方式:文档中的 Transformers 命令使用模型路径,而 Sentence Transformers 示例还可以直接加载 tencent/WeMM-Embedding-2B
  5. 考虑不受支持的音频:这些模型目前不接受音频输入,因此其 MMEB-v3 音频任务得分为零。
  6. 使用服务封装脚本以简化操作:scripts/ 中的 shell 脚本封装了文档所述的 vLLM 和 SGLang 启动配置。

性能背景

README 报告了模型在 78 个 MMEB-v2 数据集和全部 190 项 MMEB-v3 任务上的结果。在 MMEB-v2 上,2B、4B 和 9B 版本报告的平均得分分别为 77.9、79.2 和 80.6。在 MMEB-v3 上,其报告的总分分别为 56.0、58.2 和 59.5。不同任务组使用不同指标:图像和视频任务使用 Hit@1,视觉文档和文本任务使用 NDCG@5,而 agent、MCMR 和音频任务使用 Hit@1。

解读基准测试得分时,应考虑文档所述的任务覆盖范围。MMEB-v3 会将不受支持的任务计为零分,其中包括 WeMM-Embedding 不支持的音频任务。

总结

WeMM-Embedding 提供了一个模型系列,可从多种视觉和文本输入中生成归一化表示。建议先从 2B checkpoint 和项目提供的推理示例开始;需要紧凑向量时,可选择受支持的 Matryoshka 维度;需要部署服务或进行评估时,再使用 vLLM、SGLang 或 MMEB-v3 流程。除非另有说明,Tencent 编写的代码仓库内容均根据 Apache License 2.0 发布。