使用 sentence-transformers 创建文本 embedding

1. 概述

Embedding 是将一句话转换为数值向量的方法。两个意思相近的句子会产生相近的向量。这是语义搜索、RAG、文本聚类和内容推荐的基础。

在本文中，你将使用 sentence-transformers 库在本地直接运行一个来自 Hugging Face 的 embedding 模型。读完后，你将学会如何创建 embedding、计算相似度、执行语义搜索、编写测试，以及将模型封装成 HTTP 服务供 Spring Boot 应用调用。

2. 安装

sentence-transformers 是一个 Python 库，运行在 PyTorch 上。你应该创建一个独立的虚拟环境：

python -m venv .venv
source .venv/bin/activate
pip install sentence-transformers

这个命令会安装 torch、transformers 和 huggingface_hub。如果你的机器有 NVIDIA GPU，应该先安装包含 CUDA 的 torch 版本，然后再安装 sentence-transformers。

第一次加载模型时，库会将其从 Hugging Face Hub 下载到缓存目录（默认为 ~/.cache/huggingface）。之后，模型直接从缓存读取。

3. 创建第一个 embedding

我们从 all-MiniLM-L6-v2 开始。这是一个小模型（约 90 MB），速度快，适合用英文文本进行实验：

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

sentences = [
    "The order was shipped yesterday.",
    "Your package left our warehouse last night.",
    "I would like to change my password.",
]

embeddings = model.encode(sentences)

print(embeddings.shape)
print(model.get_sentence_embedding_dimension())

说明：

• SentenceTransformer(...) 从 Hub 加载模型和分词器，如果已缓存则从缓存读取。
• model.encode(sentences) 接收一个字符串列表，返回一个 numpy.ndarray。每个句子对应一行。
• get_sentence_embedding_dimension() 返回向量的维数。

输出：

(3, 384)
384

你有三个句子，每个句子被转换为一个 384 维的向量。不同的模型会有不同的维数，例如 768 或 1024。

3.1. 越南文本

all-MiniLM-L6-v2 主要在英文上进行训练。对于越南文本，你应该使用多语言模型：

model = SentenceTransformer(
    "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
)

vectors = model.encode([
    "Đơn hàng đã được giao cho đơn vị vận chuyển.",
    "Gói hàng của bạn đã rời kho tối qua.",
])
print(vectors.shape)

(2, 384)

这个模型也输出 384 维的向量。它支持约 50 种语言，包括越南语。

4. 计算相似度

向量只有在相互比较时才有用。最常见的度量是余弦相似度。从第 3 版开始，库提供了 model.similarity 方法：

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

sentences = [
    "The order was shipped yesterday.",
    "Your package left our warehouse last night.",
    "I would like to change my password.",
]
embeddings = model.encode(sentences)

scores = model.similarity(embeddings, embeddings)
print(scores)

model.similarity(a, b) 返回一个 torch.Tensor 矩阵，其中 [i][j] 位置是第 i 个句子和第 j 个句子之间的相似度。输出大致如下（由于版本不同，数值可能略有差异）：

tensor([[1.0000, 0.6123, 0.0412],
        [0.6123, 1.0000, 0.0587],
        [0.0412, 0.0587, 1.0000]])

对角线总是 1，因为每个句子与自己完全相同。两个关于送货的句子得分较高，尽管几乎没有词汇重叠。关于改密码的句子似乎毫无关联。

4.1. 归一化向量

如果你计划将向量存储到向量数据库中，请将它们归一化为长度 1：

embeddings = model.encode(sentences, normalize_embeddings=True)
score = embeddings[0] @ embeddings[1]
print(round(float(score), 4))

0.6123

当向量被归一化时，点积（@）给出正确的余弦相似度值。这个操作更便宜，许多数据库如 pgvector、Qdrant 或 Elasticsearch 都针对这个操作进行了优化。

5. 语义搜索

一个实际问题是在 FAQ 库中为客户问题找到最合适的答案。util.semantic_search 函数可以做到这一点：

from sentence_transformers import SentenceTransformer, util

model = SentenceTransformer(
    "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
)

faq = [
    "Làm sao để theo dõi đơn hàng?",
    "Chính sách hoàn tiền như thế nào?",
    "Tôi quên mật khẩu đăng nhập.",
    "Phí vận chuyển ra nước ngoài là bao nhiêu?",
]
faq_embeddings = model.encode(faq, convert_to_tensor=True, normalize_embeddings=True)

query = "Hàng của tôi đang ở đâu rồi?"
query_embedding = model.encode(query, convert_to_tensor=True, normalize_embeddings=True)

hits = util.semantic_search(query_embedding, faq_embeddings, top_k=2)[0]
for hit in hits:
    print(f"{hit['score']:.3f}  {faq[hit['corpus_id']]}")

说明：

• 你将整个 FAQ 编码一次。在真实应用中，这一步在启动时运行，或者在数据变化时运行。
• convert_to_tensor=True 将结果保存为张量，避免了 numpy 和张量之间的转换。
• semantic_search 为每个查询返回一个列表。每个元素是一个包含 corpus_id 和 score 的字典，按分数降序排列。

样本输出：

0.721  Làm sao để theo dõi đơn hàng?
0.388  Phí vận chuyển ra nước ngoài là bao nhiêu?

这种方法对于内存中的几万个文档来说已经足够。对于数百万文档，你应该切换到具有 ANN 索引的向量数据库。

注意： 某些模型如 E5 或 BGE 系列要求添加前缀，例如 "query: " 用于问题，"passage: " 用于文档。省略前缀会显著降低搜索质量。在使用前请阅读模型卡。你可以在 encode 中传递 prompt="query: "。

6. 性能优化

encode 有几个参数直接影响速度：

import torch
from sentence_transformers import SentenceTransformer

device = "cuda" if torch.cuda.is_available() else "cpu"
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device=device)

documents = [f"Customer review number {i}" for i in range(10_000)]

embeddings = model.encode(
    documents,
    batch_size=64,
    show_progress_bar=True,
    normalize_embeddings=True,
)
print(embeddings.shape, device)

Batches: 100%|██████████| 157/157 [00:09<00:00, 16.8it/s]
(10000, 384) cpu

• device 选择 CPU 或 GPU。在 GPU 上，速度可能快数十倍。
• batch_size 默认为 32。在 GPU 上，增加到 64–256 通常可以加速，只要不会耗尽 VRAM。
• 始终将整个列表传入 encode。在 for 循环中逐句调用 encode 会慢得多。

每个模型都有输入长度限制，查看 model.max_seq_length。对于 all-MiniLM-L6-v2，这个值是 256 个 token。超过这个长度的文本会被静默截断，没有警告。因此，长文档在编码前需要进行分块。

7. 为 embedding 编写测试

Embedding 是实数，所以你不能比较固定的值。应该测试属性而不是具体值。安装 pytest：

pip install pytest

# test_embeddings.py
import numpy as np
import pytest
from sentence_transformers import SentenceTransformer


@pytest.fixture(scope="module")
def model():
    return SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")


def test_embedding_has_expected_dimension(model):
    vectors = model.encode(["Order #42 is delayed"])
    assert vectors.shape == (1, 384)


def test_normalized_embedding_has_unit_length(model):
    vector = model.encode("Refund request", normalize_embeddings=True)
    assert np.linalg.norm(vector) == pytest.approx(1.0, abs=1e-5)


def test_similar_sentences_score_higher(model):
    shipping = model.encode("The order was shipped yesterday.")
    package = model.encode("Your package left our warehouse last night.")
    password = model.encode("I would like to change my password.")

    related = model.similarity(shipping, package).item()
    unrelated = model.similarity(shipping, password).item()
    assert related > unrelated

Fixture 的 scope="module" 意味着模型对整个文件只加载一次。第三个测试可以捕捉到有人不小心切换到不适合的模型时的错误。

8. 从 Spring Boot 应用中调用模型

你的大部分后端可能是用 Java 编写的。简单的方法是用 FastAPI 将模型包装成一个小的 HTTP 服务，然后从 Spring Boot 调用它。

8.1. FastAPI embedding 服务

pip install fastapi uvicorn

# app.py
from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer

app = FastAPI()
model = SentenceTransformer(
    "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
)


class EmbedRequest(BaseModel):
    texts: list[str]


@app.post("/embed")
def embed(request: EmbedRequest):
    vectors = model.encode(request.texts, normalize_embeddings=True)
    return {
        "embeddings": vectors.tolist(),
        "dimension": model.get_sentence_embedding_dimension(),
    }

uvicorn app:app --port 8000

模型在模块启动时加载一次，不是每次请求都加载。端点使用 def 而不是 async def。原因是 encode 是一个阻塞操作，FastAPI 会在单独的线程池中运行 def 函数，而不会阻塞事件循环。

8.2. Spring Boot 中的客户端

Maven 依赖：

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-web</artifactId>
</dependency>
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-test</artifactId>
    <scope>test</scope>
</dependency>

embedding.base-url=http://localhost:8000

public record EmbedRequest(List<String> texts) {}

public record EmbedResponse(List<List<Float>> embeddings, int dimension) {}

@Service
public class EmbeddingClient {

    private final RestClient restClient;

    public EmbeddingClient(RestClient.Builder builder,
                           @Value("${embedding.base-url}") String baseUrl) {
        this.restClient = builder.baseUrl(baseUrl).build();
    }

    public List<List<Float>> embed(List<String> texts) {
        EmbedResponse response = restClient.post()
                .uri("/embed")
                .body(new EmbedRequest(texts))
                .retrieve()
                .body(EmbedResponse.class);
        return response.embeddings();
    }
}

RestClient（自 Spring Framework 6.1 起提供）将 EmbedRequest 以 JSON 格式发送，然后使用 Jackson 将响应映射到 EmbedResponse 记录。你可以使用结果将其存储到 pgvector 或与其他向量进行比较。

8.3. 使用 MockRestServiceServer 测试客户端

import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.test.web.client.match.MockRestRequestMatchers.*;
import static org.springframework.test.web.client.response.MockRestResponseCreators.withSuccess;

class EmbeddingClientTest {

    @Test
    void shouldReturnEmbeddingsFromService() {
        RestClient.Builder builder = RestClient.builder();
        MockRestServiceServer server = MockRestServiceServer.bindTo(builder).build();

        server.expect(requestTo("http://localhost:8000/embed"))
              .andExpect(method(HttpMethod.POST))
              .andExpect(jsonPath("$.texts[0]").value("Đơn hàng bị trễ"))
              .andRespond(withSuccess("""
                      {"embeddings": [[0.1, 0.2, 0.3]], "dimension": 3}
                      """, MediaType.APPLICATION_JSON));

        EmbeddingClient client = new EmbeddingClient(builder, "http://localhost:8000");

        List<List<Float>> result = client.embed(List.of("Đơn hàng bị trễ"));

        assertThat(result).hasSize(1);
        assertThat(result.getFirst()).containsExactly(0.1f, 0.2f, 0.3f);
        server.verify();
    }
}

MockRestServiceServer.bindTo(builder) 必须在 builder 创建 RestClient 之前被调用。在这个测试中，这总是成立的，因为客户端仅在构造函数内调用 build()。这样，测试可以运行而不需要真正的 Python 服务。

9. 常见错误

• 混合来自多个模型的向量。 两个不同模型的向量处于不同的空间中，即使维数相同。改变模型意味着需要重新编码所有数据。
• 在一方忘记归一化。 如果文档被归一化但查询没有，点积结果会出错。对两方使用相同的 encode 配置。
• 文本太长。 超过 max_seq_length 的部分会被截断而没有警告。在编码前分块文档。
• 在每个请求中加载模型。 这需要几秒钟和大量 RAM。仅在启动时加载一次。
• 容器没有网络。 生产环境通常会阻止访问 Hub。将模型预加载到镜像中，然后设置 HF_HUB_OFFLINE=1。

注意： Embedding 不能替代精确搜索。对于订单号、电话号码或 SKU，全文搜索或常规 SQL 查询提供更好的结果。在实践中，结合两者（混合搜索）通常最有效。

10. 结论

你已经使用 sentence-transformers 用短短几行代码创建了 embedding。你还计算了余弦相似度，为 FAQ 库构建了语义搜索，并通过调整 batch_size 和 device 来加速运行。此外，你还学会了如何通过属性测试 embedding，以及如何将模型包装成 HTTP 服务供 Spring Boot 通过 RestClient 调用。

下一个合理步骤是将向量存储到向量数据库如 pgvector，然后在 RAG 管道中将其用作检索部分。