端侧 Embedding 模型实战：为本地 RAG 选对编码器

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

# 用 sentence-transformers 测试推理速度
from sentence_transformers import SentenceTransformer
import time

models = [
    "BAAI/bge-small-zh-v1.5",
    "shibing624/text2vec-base-chinese",
    "sentence-transformers/all-MiniLM-L6-v2",
    "Alibaba-NLP/gte-small",
]

sentences = [
    "Python 中的列表推导式比 for 循环快多少？",
    "如何用 Docker 优化前端构建流程？",
    "端侧部署大模型的常见量化方案对比",
    "用 ChromaDB 做本地 RAG 的完整步骤",
]

for model_name in models:
    model = SentenceTransformer(model_name, device="cpu")
    # 预热
    _ = model.encode(sentences)
    # 实测 100 轮
    start = time.time()
    for _ in range(100):
        _ = model.encode(sentences, show_progress_bar=False)
    elapsed = time.time() - start
    print(f"{model_name:50s} 100 轮耗时: {elapsed:.2f}s  "
          f"单次: {elapsed/100/len(sentences)*1000:.1f}ms/条")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

import psutil, os
from sentence_transformers import SentenceTransformer

def get_model_mb(model_name):
    model = SentenceTransformer(model_name, device="cpu")
    # 加载后的内存增量
    before = psutil.Process(os.getpid()).memory_info().rss
    # 编码一次触发实际加载
    _ = model.encode(["test"])
    after = psutil.Process(os.getpid()).memory_info().rss
    return (after - before) / 1024 / 1024

for name in models:
    mb = get_model_mb(name)
    print(f"{name:50s} 内存占用: {mb:.0f} MB")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

from sentence_transformers import SentenceTransformer, util

# 构造 100 条问答对
corpus = {
    "什么是 Python 的 GIL？": "GIL（全局解释器锁）是 CPython 中的一个互斥锁，"
                              "限制同一时刻只有一个线程执行 Python 字节码。",
    "Docker 和虚拟机有什么区别？": "Docker 共享宿主机内核，启动毫秒级；"
                                   "虚拟机包含完整操作系统，启动分钟级。",
    # ... 更多问答对
}

model = SentenceTransformer("BAAI/bge-small-zh-v1.5")
corpus_texts = list(corpus.keys())
corpus_emb = model.encode(corpus_texts, convert_to_tensor=True)

def search(query, top_k=5):
    query_emb = model.encode(query, convert_to_tensor=True)
    scores = util.cos_sim(query_emb, corpus_emb)[0]
    top_indices = scores.topk(top_k).indices.tolist()
    return [corpus_texts[i] for i in top_indices]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.config import Settings

# 初始化端侧 RAG
model = SentenceTransformer("BAAI/bge-small-zh-v1.5", device="cpu")
client = chromadb.PersistentClient(
    path="./local_rag_db",
    settings=Settings(anonymized_telemetry=False)
)
collection = client.get_or_create_collection(
    name="docs",
    embedding_function=model.encode  # 直接传入
)

1
2
3
4
5
6
7
8
9

# 下载 GGUF 格式的 embedding 模型
# 例如：bge-small-en-v1.5-ggml
wget https://huggingface.co/ChristianAzinn/bge-small-en-v1.5-ggml/resolve/main/bge-small-en-v1.5-Q4_K_M.gguf

# 嵌入推理
./llama-embedding \
  --model bge-small-en-v1.5-Q4_K_M.gguf \
  --prompt "端侧部署的量化方案有哪些？" \
  --embd-normalize 2

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

import onnxruntime as ort
from transformers import AutoTokenizer
import numpy as np

model_id = "BAAI/bge-small-zh-v1.5"
# 先导出为 ONNX
# python -m optimum.exporters.onnx --model BAAI/bge-small-zh-v1.5 bge_onnx/

session = ort.InferenceSession("bge_onnx/model.onnx")
tokenizer = AutoTokenizer.from_pretrained(model_id)

def embed_onnx(text):
    inputs = tokenizer(text, return_tensors="np", padding=True, truncation=True)
    outputs = session.run(None, {
        "input_ids": inputs["input_ids"],
        "attention_mask": inputs["attention_mask"]
    })
    # 取 CLS token 作为 sentence embedding
    embedding = outputs[0][:, 0, :].squeeze()
    # L2 归一化
    return embedding / np.linalg.norm(embedding)

1
2
3
4
5

# 错误：向量未归一化，余弦相似度结果不准
embeddings = model.encode(texts)

# 正确：显式归一化
embeddings = model.encode(texts, normalize_embeddings=True)

1
2
3
4
5

# 长文档会静默截断，丢失尾部信息
model.max_seq_length = 512  # 默认 512

# 用 bge-m3 可以处理 8192 token
# 或用 jina-embeddings 处理 8192 token