开源文本嵌入模型的高效应用

原创

用户11764306

发布于 2026-05-26 12:16:35

使用开源模型实现更快、更便宜的文本嵌入

嵌入是处理文本的强大工具。通过将文本“嵌入”到向量中，可以将其含义编码为一种表示形式，更易于用于语义搜索、聚类和分类等任务。

本指南将演示如何在Replicate平台上使用BAAI/bge-large-en-v1.5模型生成文本嵌入。“BAAI通用嵌入”（BGE）模型套件由某机构发布，是开源的。

截至2023年10月，本文使用的大型BGE模型是当前最先进的开源文本嵌入模型。它在MTEB排行榜上排名高于某机构的嵌入模型，且在Replicate上大规模运行文本嵌入的成本便宜4倍。

环境准备

需要以下条件：

一个Replicate平台账户（用于运行BGE模型）
一个Python环境（或使用Google Colab笔记本）

安装依赖

首先安装以下依赖库：

pip install replicate
pip install transformers sentencepiece
pip install datasets py7zr scikit-learn

身份验证

从Replicate账户获取API令牌，并设置为环境变量：

export REPLICATE_API_TOKEN=...

从文本列表生成嵌入

使用replicate库运行嵌入模型：

import json
import replicate

texts = [
    "the happy cat",
    "the quick brown fox jumps over the lazy dog",
    "lorem ipsum dolor sit amet",
    "this is a test",
]

output = replicate.run(
    "nateraw/bge-large-en-v1.5:9cf9f015a9cb9c61d1a2610659cdac4a4ca222f2d3707a68517b18c198a9add1",
    input={"texts": json.dumps(texts)}
)
print(output)

输出将是每个文本对应的嵌入向量列表。

从JSONL文件生成嵌入

JSONL（JSON Lines）是一种基于文本、行分隔格式的结构化数据存储格式。每行是一个独立的JSON对象。

示例文件 dummy_example.jsonl：

{"text": "the happy cat"}
{"text": "the quick brown fox jumps over the lazy dog"}
{"text": "lorem ipsum dolor sit amet"}
{"text": "this is a test"}

通过指定path输入运行模型：

output = replicate.run(
    "nateraw/bge-large-en-v1.5:9cf9f015a9cb9c61d1a2610659cdac4a4ca222f2d3707a68517b18c198a9add1",
    input={"path": open("dummy_example.jsonl", "rb")}
)
len(output)  # 输出: 4

实际案例：嵌入SAMSum数据集

SAMSum数据集包含约1.4万条对话示例及其人工标注摘要，常用于训练和评估语言模型。

以下代码将整个SAMSum数据集编码，转换为JSONL文件，然后运行BGE模型生成文本嵌入：

from pathlib import Path
from datasets import load_dataset

dataset_name = "samsum"
text_field = "dialogue"
outfile_name = "samsum_dialogue.jsonl"

ds = load_dataset(dataset_name, split='train')
ds = ds.remove_columns([x for x in ds.column_names if x != text_field])
ds = ds.rename_column(text_field, "text")
ds.to_json(outfile_name)

嵌入整个数据集，指定convert_to_numpy=True以获得numpy数组格式：

import time

start = time.time()
output = replicate.run(
    "nateraw/bge-large-en-v1.5:9cf9f015a9cb9c61d1a2610659cdac4a4ca222f2d3707a68517b18c198a9add1",
    input=dict(
        path=open(outfile_name, "rb"),
        convert_to_numpy=True,
        batch_size=64
    )
)
time_to_embed = time.time() - start
print(f"耗时: {time_to_embed:.2f} 秒")
print("输出地址:", output)
# 输出: 耗时: 65.51 秒.

加载生成的嵌入：

import requests
from io import BytesIO
import numpy as np

embeds = np.load(BytesIO(requests.get(output).content))
embeds.shape  # 输出: (14732, 1024)

成本对比

某机构的Ada v2模型价格为$0.0001 / 1K tokens。

Replicate平台上运行A40（Large）实例的成本为$0.000725/秒。

对包含5,120,000个token的数据集进行基准测试：

total_tokens = 5120000

# 某机构API成本
openai_cost = 0.0001  # 每1k tokens
openai_price = total_tokens / 1000 * openai_cost
print(f"某机构价格: ${openai_price:.3f} USD")
# 某机构价格: $0.512 USD

# Replicate平台成本
time_to_embed = 151.92  # 秒
replicate_price = time_to_embed * 0.000725
print(f"某平台价格: ${replicate_price:.3f}")
# 某平台价格: $0.110

某平台的价格比某机构便宜4倍以上，且使用的模型在MTEB排行榜上排名更高。FINISHED

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

网络安全