文本 embedding(BEG模型)
更新时间：2026-06-15
简介
基于 BGE 系列的文本嵌入模型，支持稠密向量生成
功能描述
- 支持 BGE-M3、BGE-Large-zh-v1.5、BGE-Large-en-v1.5、BGE-Multilingual-Gemma2 等模型
- 每个模型有不同的输入提示，用户可以根据需要选择不同的模型
- 硬件加速**：支持
FP16量化与 GPU 并行计算
算子参数
输入
| 输入 | 含义 |
|---|---|
| texts | 包含待处理文本的数组，元素类型为str。 |
输出
| 输出 | 含义 |
|---|---|
| result | 处理后的数组，包含每个文本对应的稠密嵌入向量 |
参数
| 参数名称 | 类型 | 默认值 | 描述 |
|---|---|---|---|
| dtype | str | 'float32' | 模型精度，支持 float32 和 float16 可选值：["float32", "float16"] 默认值："float32" |
| batch_size | int | 512 | 模型推理时的批处理大小 默认值：512 |
| model_path | str | '/opt/aihc/models' | 模型文件所在的路径 默认值："/opt/aihc/models" |
| model_name | str | 'BAAI/bge-m3' | 模型名称 可选值：["BAAI/bge-m3", "BAAI/bge-large-zh-v1.5", "BAAI/bge-large-en-v1.5", "BAAI/bge-multilingual-gemma2"] 默认值："BAAI/bge-m3" |
| rank | int | None | GPU 编号 默认值：None |
调用示例
Python
1from __future__ import annotations
2
3import os
4
5import daft
6from daft import col
7
8from daft.aihc.common.udf import aihc_udf
9from daft.aihc.functions.text.embedding.bge_embedding import BgeEmbedding
10
11if __name__ == "__main__":
12 if os.getenv("DAFT_RUNNER", "native") == "ray":
13 import ray
14 ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
15 daft.set_runner_ray()
16 daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
17
18 # TODO: 根据实际场景准备样本数据
19 samples = {"texts": [...]}
20 ds = daft.from_pydict(samples)
21 constructor_kwargs = {
22 "dtype": 'float32',
23 "batch_size": 512,
24 "model_path": '/opt/aihc/models',
25 "model_name": 'BAAI/bge-m3',
26 "rank": None,
27 }
28 ds = ds.with_column(
29 "result",
30 aihc_udf(
31 BgeEmbedding,
32 construct_args=constructor_kwargs,
33 num_cpus=1,
34 concurrency=4,
35 batch_size=8,
36 )(col("texts")),
37 )
38 ds.show()评价此篇文章