第4章 RAG — 02 Embedding 与语义搜索#
Embedding（嵌入）把文字转成向量，语义相似的文字在向量空间中距离近。这是 RAG 检索的核心。
本节将从零实现一个完整的语义搜索系统，涵盖：
步骤 |
内容 |
|---|---|
1 |
调用 Embedding 模型，获取文本向量 |
2 |
用余弦相似度衡量语义距离 |
3 |
构建简单向量索引，实现 Top-K 检索 |
4 |
语义搜索 vs 关键词搜索的对比 |
5 |
Embedding 缓存，避免重复计算 |
前置要求：已配置
EMBED_MODEL环境变量（或使用默认的openai/text-embedding-3-small）
import json
import os
import hashlib
import time
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import litellm
from dotenv import load_dotenv
load_dotenv()
# Embedding 模型配置
EMBED_MODEL = os.getenv("EMBED_MODEL", "openai/text-embedding-3-small")
print(f"Embedding 模型: {EMBED_MODEL}")
print(f"NumPy 版本: {np.__version__}")
Embedding 模型: openai/text-embedding-3-small
NumPy 版本: 2.4.3
第一节：生成 Embedding#
通过 litellm.embedding() 调用任意 Embedding 模型（OpenAI、Cohere、本地模型等），
返回一个高维浮点数向量。向量的每个维度代表文本的某种语义特征。
text-embedding-3-small：1536 维，性价比高text-embedding-3-large：3072 维，精度更高
def get_embedding(text, model=EMBED_MODEL):
"""
获取文本的 Embedding 向量
:param text: 输入文本字符串
:param model: Embedding 模型名称
:return: numpy.ndarray，形状为 (dim,)
"""
response = litellm.embedding(model=model, input=[text])
vector = np.array(response.data[0]["embedding"])
return vector
# 演示：对一句话生成 Embedding
sample_sentence = "人工智能正在改变世界"
print(f"输入文本: {sample_sentence}")
print("正在调用 Embedding 模型...")
vec = get_embedding(sample_sentence)
print(f"\n向量维度: {vec.shape[0]}")
print(f"向量数据类型: {vec.dtype}")
print(f"向量 L2 范数: {np.linalg.norm(vec):.4f}")
print(f"\n前 8 个维度的值:")
print(np.round(vec[:8], 4))
输入文本: 人工智能正在改变世界
正在调用 Embedding 模型...
向量维度: 1536
向量数据类型: float64
向量 L2 范数: 1.0002
前 8 个维度的值:
[ 0.0364 0.0006 0.0407 0.0559 0.0335 -0.0342 -0.0125 0.0756]
第二节：余弦相似度（Cosine Similarity）#
余弦相似度衡量两个向量夹角的余弦值，取值范围 [-1, 1]：
1.0：完全相同0.0：无关-1.0：完全相反
对5个句子（涵盖相似和不同主题）计算两两余弦相似度矩阵，直观看出语义距离。
# 5 个句子：混合相似和不同主题
sentences = [
"深度学习是机器学习的一个重要分支", # 0 - AI/ML
"神经网络可以从大量数据中自动学习特征", # 1 - AI/ML（与0相似）
"Transformer 模型彻底改变了自然语言处理", # 2 - NLP
"今天天气晴朗，非常适合户外运动", # 3 - 天气（不同主题）
"篮球比赛非常精彩，球迷们热情高涨", # 4 - 体育（不同主题）
]
print("正在为 5 个句子生成 Embedding...")
vectors = np.array([get_embedding(s) for s in sentences])
print(f"Embedding 矩阵形状: {vectors.shape} (句子数 × 向量维度)")
# 计算余弦相似度矩阵
sim_matrix = cosine_similarity(vectors) # shape: (5, 5)
# 格式化打印矩阵
print("\n余弦相似度矩阵（行=句子索引，列=句子索引）:")
print("\n ", end="")
for i in range(len(sentences)):
print(f" S{i} ", end="")
print()
print("-" * 38)
for i, row in enumerate(sim_matrix):
print(f"S{i} | ", end="")
for val in row:
print(f"{val:.3f} ", end="")
suffix = f" \u2190 {sentences[i][:18]}..." if len(sentences[i]) > 18 else f" \u2190 {sentences[i]}"
print(suffix)
# 找出相似度最高的非对角线对
print("\n最相似的句子对（排除自身）:")
pairs = []
for i in range(len(sentences)):
for j in range(i+1, len(sentences)):
pairs.append((sim_matrix[i][j], i, j))
pairs.sort(reverse=True)
for sim, i, j in pairs[:3]:
print(f" [{sim:.3f}] S{i} vs S{j}: \"{sentences[i][:15]}...\" \u2194 \"{sentences[j][:15]}...\"")
正在为 5 个句子生成 Embedding...
Embedding 矩阵形状: (5, 1536) (句子数 × 向量维度)
余弦相似度矩阵（行=句子索引，列=句子索引）:
S0 S1 S2 S3 S4
--------------------------------------
S0 | 1.000 0.423 0.373 0.109 0.135 ← 深度学习是机器学习的一个重要分支
S1 | 0.423 1.000 0.349 0.080 0.151 ← 神经网络可以从大量数据中自动学习特征
S2 | 0.373 0.349 1.000 0.106 0.073 ← Transformer 模型彻底改变...
S3 | 0.109 0.080 0.106 1.000 0.346 ← 今天天气晴朗，非常适合户外运动
S4 | 0.135 0.151 0.073 0.346 1.000 ← 篮球比赛非常精彩，球迷们热情高涨
最相似的句子对（排除自身）:
[0.423] S0 vs S1: "深度学习是机器学习的一个重要分..." ↔ "神经网络可以从大量数据中自动学..."
[0.373] S0 vs S2: "深度学习是机器学习的一个重要分..." ↔ "Transformer 模型彻..."
[0.349] S1 vs S2: "神经网络可以从大量数据中自动学..." ↔ "Transformer 模型彻..."
第三节：构建简单向量索引（SimpleVectorIndex）#
向量索引的核心操作：
add(texts)：批量添加文本，存储其 Embedding 向量search(query, top_k)：对查询文本也生成 Embedding，计算与所有存储向量的余弦相似度，返回 Top-K
这是生产级向量数据库（Faiss、Chroma、Weaviate）的简化原理。
class SimpleVectorIndex:
"""
基于 NumPy 的简单向量索引
支持批量添加文本和 Top-K 语义搜索
"""
def __init__(self, embed_model=EMBED_MODEL):
self.embed_model = embed_model
self.texts = [] # 原始文本列表
self.vectors = None # 向量矩阵 (N, dim)
def add(self, texts):
"""
批量添加文本到索引
:param texts: List[str]
"""
print(f"正在为 {len(texts)} 条文本生成 Embedding...")
new_vectors = []
for text in texts:
vec = get_embedding(text, model=self.embed_model)
new_vectors.append(vec)
new_matrix = np.array(new_vectors)
if self.vectors is None:
self.vectors = new_matrix
else:
self.vectors = np.vstack([self.vectors, new_matrix])
self.texts.extend(texts)
print(f"索引中共有 {len(self.texts)} 条文本")
def search(self, query, top_k=3):
"""
语义搜索
:param query: 查询文本
:param top_k: 返回最相似的 K 条
:return: List of (score, text)
"""
if self.vectors is None:
raise ValueError("索引为空，请先调用 add()")
query_vec = get_embedding(query, model=self.embed_model)
query_vec = query_vec.reshape(1, -1) # (1, dim)
scores = cosine_similarity(query_vec, self.vectors)[0] # (N,)
top_indices = np.argsort(scores)[::-1][:top_k]
return [(float(scores[i]), self.texts[i]) for i in top_indices]
# ---- 演示 ----
corpus = [
"机器学习是让计算机从数据中自动学习的技术",
"深度神经网络由多层感知机堆叠而成",
"卷积神经网络擅长处理图像识别任务",
"自然语言处理让计算机理解人类语言",
"大语言模型通过海量文本预训练获得语言能力",
"强化学习通过奖励信号训练智能体做决策",
"Python 是数据科学最流行的编程语言",
"向量数据库用于存储和检索高维 Embedding",
"今天股市大涨，沪指突破3500点",
"足球世界杯吸引了全球数十亿观众观看",
]
index = SimpleVectorIndex()
index.add(corpus)
# 搜索
query = "机器学习算法"
print(f"\n查询: '{query}'")
print("-" * 50)
results = index.search(query, top_k=3)
for rank, (score, text) in enumerate(results, 1):
print(f"Top-{rank} [相似度={score:.4f}]: {text}")
正在为 10 条文本生成 Embedding...
索引中共有 10 条文本
查询: '机器学习算法'
--------------------------------------------------
Top-1 [相似度=0.6591]: 机器学习是让计算机从数据中自动学习的技术
Top-2 [相似度=0.5035]: 自然语言处理让计算机理解人类语言
Top-3 [相似度=0.4257]: 强化学习通过奖励信号训练智能体做决策
第四节：语义搜索 vs 关键词搜索#
关键词搜索（字符串匹配）：只能找到包含查询词的文本，无法理解同义词或相关概念。
语义搜索（Embedding 相似度）：理解查询意图，即使没有完全匹配的词汇也能找到相关内容。
用同一个语料库和查询词，并排展示两种方法的差异。
def keyword_search(query, texts, top_k=3):
"""
基础关键词搜索：统计查询词中各字在文本中出现的次数作为分数
:param query: 查询字符串
:param texts: 语料库 List[str]
:param top_k: 返回 Top-K
:return: List of (score, text)
"""
scored = []
query_chars = set(query) # 简单用字符级匹配
for text in texts:
# 计算查询词中有多少字出现在文本里
score = sum(1 for ch in query_chars if ch in text)
scored.append((score, text))
scored.sort(key=lambda x: -x[0])
# 过滤掉分数为 0 的结果
results = [(s, t) for s, t in scored[:top_k] if s > 0]
return results
# 同一语料库（已在上一节的 index 中）
query = "深度学习模型训练技巧"
print(f"查询: '{query}'")
print("=" * 60)
# --- 关键词搜索 ---
kw_results = keyword_search(query, corpus, top_k=3)
print("\n[关键词搜索结果]")
if kw_results:
for rank, (score, text) in enumerate(kw_results, 1):
print(f" Top-{rank} [匹配字数={score}]: {text}")
else:
print(" 无匹配结果")
# --- 语义搜索 ---
sem_results = index.search(query, top_k=3)
print("\n[语义搜索结果]")
for rank, (score, text) in enumerate(sem_results, 1):
print(f" Top-{rank} [相似度={score:.4f}]: {text}")
# 差异分析
print("\n" + "=" * 60)
print("差异分析:")
kw_texts = {t for _, t in kw_results}
sem_texts = {t for _, t in sem_results}
only_semantic = sem_texts - kw_texts
if only_semantic:
print("语义搜索找到但关键词搜索未找到的结果:")
for t in only_semantic:
print(f" \u2192 {t}")
else:
print("两种方法结果相同")
查询: '深度学习模型训练技巧'
============================================================
[关键词搜索结果]
Top-1 [匹配字数=4]: 大语言模型通过海量文本预训练获得语言能力
Top-2 [匹配字数=4]: 强化学习通过奖励信号训练智能体做决策
Top-3 [匹配字数=3]: 机器学习是让计算机从数据中自动学习的技术
[语义搜索结果]
Top-1 [相似度=0.5022]: 大语言模型通过海量文本预训练获得语言能力
Top-2 [相似度=0.4782]: 机器学习是让计算机从数据中自动学习的技术
Top-3 [相似度=0.4758]: 强化学习通过奖励信号训练智能体做决策
============================================================
差异分析:
两种方法结果相同
第五节：Embedding 缓存（Caching）#
Embedding API 调用有网络延迟和费用。对于相同的文本，多次调用会返回完全相同的向量。
缓存策略：以文本内容的 SHA256 哈希为键，将向量存储在内存字典中。
首次请求：调用 API，存入缓存（Cache Miss）
再次请求：直接从缓存读取，无需 API 调用（Cache Hit）
class CachedEmbedder:
"""
带缓存的 Embedding 客户端
使用文本 SHA256 哈希作为缓存键，避免重复 API 调用
"""
def __init__(self, model=EMBED_MODEL):
self.model = model
self._cache = {} # {text_hash: np.ndarray}
self.hits = 0 # 缓存命中次数
self.misses = 0 # 缓存未命中次数
def _hash(self, text):
"""计算文本的 SHA256 哈希"""
return hashlib.sha256(text.encode("utf-8")).hexdigest()
def embed(self, text):
"""
获取 Embedding，优先从缓存读取
:param text: 输入文本
:return: np.ndarray
"""
key = self._hash(text)
if key in self._cache:
self.hits += 1
return self._cache[key]
# Cache miss：调用 API
self.misses += 1
vector = get_embedding(text, model=self.model)
self._cache[key] = vector
return vector
def stats(self):
"""返回缓存统计信息"""
total = self.hits + self.misses
hit_rate = self.hits / total * 100 if total > 0 else 0
return {
"total_calls": total,
"cache_hits": self.hits,
"cache_misses": self.misses,
"hit_rate": f"{hit_rate:.1f}%",
"cached_vectors": len(self._cache),
}
# ---- 演示 ----
embedder = CachedEmbedder()
texts_to_embed = [
"人工智能正在快速发展",
"机器学习改变了软件开发方式",
"人工智能正在快速发展", # 重复文本，应命中缓存
"深度学习需要大量训练数据",
"机器学习改变了软件开发方式", # 重复文本，应命中缓存
"人工智能正在快速发展", # 再次重复
]
print(f"共 {len(texts_to_embed)} 次 embed 调用（含重复文本）\n")
print(f"{'调用#':<5} {'Cache':<12} {'文本':<35}")
print("-" * 55)
for i, text in enumerate(texts_to_embed, 1):
key = embedder._hash(text)
is_hit = key in embedder._cache
status = "HIT \u2713" if is_hit else "MISS \u2192API"
t0 = time.time()
vec = embedder.embed(text)
elapsed = time.time() - t0
print(f"#{i:<4} {status:<12} {text:<35} ({elapsed*1000:.0f}ms)")
# 打印统计
print("\n" + "=" * 55)
print("缓存统计:")
stats = embedder.stats()
for k, v in stats.items():
print(f" {k:<20}: {v}")
共 6 次 embed 调用（含重复文本）
调用# Cache 文本
-------------------------------------------------------
#1 MISS →API 人工智能正在快速发展 (305ms)
#2 MISS →API 机器学习改变了软件开发方式 (241ms)
#3 HIT ✓ 人工智能正在快速发展 (0ms)
#4 MISS →API 深度学习需要大量训练数据 (188ms)
#5 HIT ✓ 机器学习改变了软件开发方式 (0ms)
#6 HIT ✓ 人工智能正在快速发展 (0ms)
=======================================================
缓存统计:
total_calls : 6
cache_hits : 3
cache_misses : 3
hit_rate : 50.0%
cached_vectors : 3
扩展：持久化缓存#
上面的内存缓存在程序重启后会丢失。可以将缓存序列化到磁盘（JSON / SQLite / NumPy）， 下次启动时加载，实现跨会话复用。
import tempfile
import pathlib
class PersistentEmbedder(CachedEmbedder):
"""
持久化 Embedding 缓存（存储为 .npz 文件）
- save()：将缓存写入磁盘
- load()：从磁盘加载缓存
"""
def __init__(self, cache_path="embed_cache.npz", **kwargs):
super().__init__(**kwargs)
self.cache_path = pathlib.Path(cache_path)
self.load() # 启动时自动加载已有缓存
def save(self):
"""将内存缓存保存到 .npz 文件"""
if not self._cache:
print("缓存为空，跳过保存")
return
# npz 文件的键名不能包含特殊字符，使用索引映射
keys = list(self._cache.keys())
arrays = {f"v{i}": self._cache[k] for i, k in enumerate(keys)}
# 将 key 列表也保存进去
arrays["__keys__"] = np.array(keys)
np.savez(self.cache_path, **arrays)
print(f"缓存已保存: {self.cache_path} ({len(keys)} 条向量)")
def load(self):
"""从 .npz 文件加载缓存"""
if not self.cache_path.exists():
print(f"缓存文件不存在，从空缓存开始: {self.cache_path}")
return
data = np.load(self.cache_path, allow_pickle=True)
keys = data["__keys__"].tolist()
for i, k in enumerate(keys):
self._cache[k] = data[f"v{i}"]
print(f"缓存已加载: {self.cache_path} ({len(keys)} 条向量)")
# 演示：保存与加载
with tempfile.TemporaryDirectory() as tmpdir:
cache_file = os.path.join(tmpdir, "embed_cache.npz")
# 第一次：生成并保存
print("=== 第一次运行（写入缓存）===")
p1 = PersistentEmbedder(cache_path=cache_file)
p1.embed("向量数据库是 RAG 的核心存储")
p1.embed("Embedding 模型将文本映射到高维空间")
p1.save()
print(f"缓存统计: {p1.stats()}\n")
# 第二次：加载缓存，无需重新调用 API
print("=== 第二次运行（读取缓存）===")
p2 = PersistentEmbedder(cache_path=cache_file)
p2.embed("向量数据库是 RAG 的核心存储") # 应命中缓存
p2.embed("Embedding 模型将文本映射到高维空间") # 应命中缓存
p2.embed("全新的文本，缓存中没有") # 未命中，调用 API
print(f"缓存统计: {p2.stats()}")
=== 第一次运行（写入缓存）===
缓存文件不存在，从空缓存开始: /var/folders/km/c2z87ytx2fvbh8_v94cn7l4w0000gp/T/tmp13h9651x/embed_cache.npz
缓存已保存: /var/folders/km/c2z87ytx2fvbh8_v94cn7l4w0000gp/T/tmp13h9651x/embed_cache.npz (2 条向量)
缓存统计: {'total_calls': 2, 'cache_hits': 0, 'cache_misses': 2, 'hit_rate': '0.0%', 'cached_vectors': 2}
=== 第二次运行（读取缓存）===
缓存已加载: /var/folders/km/c2z87ytx2fvbh8_v94cn7l4w0000gp/T/tmp13h9651x/embed_cache.npz (2 条向量)
缓存统计: {'total_calls': 3, 'cache_hits': 2, 'cache_misses': 1, 'hit_rate': '66.7%', 'cached_vectors': 3}
小结#
Embedding 是 RAG 检索的骨干技术：
概念 |
要点 |
|---|---|
Embedding |
文本 → 高维浮点向量，语义相近则向量相近 |
余弦相似度 |
衡量两个向量的方向相似性，与长度无关 |
向量索引 |
存储文本向量，支持快速 Top-K 相似度检索 |
语义 vs 关键词 |
语义搜索理解意图，可跨越词汇差异找到相关内容 |
缓存 |
相同文本的 Embedding 结果固定，缓存可大幅降低成本和延迟 |
下一步：
使用生产级向量数据库（Chroma、Faiss、Weaviate、Pinecone）替换 SimpleVectorIndex
学习 Hybrid Search（向量搜索 + BM25 关键词搜索的融合）
探索 Reranker 模型对检索结果的精排