Milvus上的高级RAG技术:混合搜索、HyDE、查询重写与重排序。
逐步构建Milvus RAG流程:密集向量与BM25向量、RRF混合搜索、元数据过滤、查询重写、HyDE、分解与交叉编码器重排序。
来源: Beyond Semantic Search: The Complete Guide to Advanced RAG with Milvus | Sagar Patil
一个概念验证版的RAG系统只需几分钟时间:先将PDF分割成块,再将这些块输入嵌入模型,最后将匹配结果传递给大语言模型。但实际应用中,仅依赖向量搜索的方案就会遇到瓶颈。真实用户不会像结构化数据引擎那样精确地表述查询,他们使用的措辞往往比较模糊,会遗漏关键词,还将多个问题合并在一起,而单纯的语义检索很难捕捉到这些信息。当检索失败时,大语言模型不得不对无关内容进行总结,这往往会导致细微的幻觉或毫无用处的答案。
要缩小基础语义搜索与多阶段检索流程之间的差距,需要多层解决方案。下面的实现使用了Milvus技术:首先进行密集向量检索,再加入混合搜索方式,随后组合四种技术(元数据过滤、查询重写、HyDE及查询分解),最后通过交叉编码器进行重新排序。每个步骤都配有可运行的代码。
为何会有 RAG
像 GPT-4 这样的模型仅能从知识截止日期之前的通用互联网数据中学习,无法获取私人文档、内部维基及公司政策等信息。若向这类模型询问公司的休假政策,它要么表示不知道,要么编造出一个看似合理但实际上错误的答案。
RAG(检索增强生成)将任务分为两个步骤:
- 通过语义搜索从文档中检索最相关的段落。
- 将这些段落作为上下文传递给大语言模型,从而生成答案。
整个流程如下所示:
User Question
│
▼
[Embed the question] → query vector
│
▼
[Search Vector DB] → top-K relevant document chunks
│
▼
[LLM prompt: "Given these passages, answer: {question}"]
│
▼
Accurate, Grounded Answer
最终答案的质量几乎完全取决于检索质量,因此做好检索才是关键。
向量数据库的作用
传统数据库比较的是精确值,因此即使“vacation entitlement”和“medical absence”的含义相同,WHERE text LIKE '%leave%'查询也会忽略这些文档。而向量数据库则将每篇文档存储为高维数值向量(即嵌入向量),并通过语义相似度进行检索,即判断查询的含义与各文档片段含义的接近程度。基于以下原因,Milvus非常适合这一角色:
- 它是开源的,可以自行托管,从而保障数据隐私。
- Milvus Lite仅需一个
.db文件即可运行,开发无需服务器。
密集型与稀疏型嵌入
该流程使用两种类型的向量,因此在编写代码之前了解它们很有必要。密集型嵌入是由神经网络生成的固定长度向量,例如包含1024个数值;所有维度都带有值,训练过程会让模型将含义相似的句子放在一起:
"sick leave policy" → [0.12, -0.87, 0.34, 0.56, ...] (1024 numbers)
"medical absence entitlement" → [0.13, -0.85, 0.31, 0.54, ...] ← very close
"quarterly revenue target" → [0.91, 0.23, -0.67, 0.02, ...] ← far away
密集向量能够很好地捕捉语义和同义词。OpenAI的text-embedding-3-large等模型就能生成这类向量。而稀疏嵌入则是一个巨大的向量,其中大部分元素为零。每个非零值对应一个词汇表中的单词,其数值代表该词的重要性,即类似TF-IDF的权重:
"sick leave policy" → {word_index_for_"sick": 0.82, word_index_for_"leave": 0.91, ...}
BM25是一种经典的稀疏算法,擅长进行精确的关键词匹配。它能够识别专有名词、产品名称和技术术语,而这些内容往往因为没有相近的语义关联而被密集搜索忽略。
混合搜索则结合了这两种方法的优势,充分利用各自的优点。
设置与配置
首先安装依赖项:
pip install --upgrade pymilvus
pip install "pymilvus[model]"
pip install sentence-transformers
pip install langchain-text-splitters
pip install langchain-openai
pip install langchain-community
pip install scipy
pip install nltk
接下来是导入内容:
import uuid
from tqdm import tqdm
from pymilvus import (
MilvusClient, DataType,
AnnSearchRequest, RRFRanker
)
from pymilvus.model.sparse import BM25EmbeddingFunction
from pymilvus.model.sparse.bm25.tokenizers import build_default_analyzer
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
import scipy.sparse as sp
import re, json
import nltk
nltk.download('stopwords')
将所有可调整的参数集中放在文件的顶部:
PDF_PATH = "./data/sample_employee_handbook.pdf" # path of you document
COLLECTION_NAME = "rag_documents_hybrid"
MILVUS_DB_PATH = "./db/milvus_demo.db"
API_KEY = "sk-..."
EMBEDDING_MODEL = "text-embedding-3-large"
EMBEDDING_DIM = 1024
CHUNK_SIZE = 500
CHUNK_OVERLAP = 100
TOP_K = 5
关于分块大小:500个字符且重叠100个字符的设置适合用作简单文档的初始方案。叙事性文本更适合800到1000个字符的大分块,而内容密集的技术文档则适合200到300个字符的小分块。重叠设置能确保被分块边界截断的句子仍可被找到。
索引流程
索引操作会执行一次,或在文档发生变化时再次执行。它将原始文档转换为向量并加载到Milvus中,具体路径为:
PDF → Pages → Chunks → Dense Embeddings
→ Sparse Embeddings
→ Milvus Collection
首先初始化模型并建立连接:
# Dense embedding model here we'll be using OpenAI's embedding model
embedding_obj = OpenAIEmbeddings(
model=EMBEDDING_MODEL,
api_key=API_KEY,
dimensions=EMBEDDING_DIM
)
# Milvus Lite - single file, no server needed
client = MilvusClient(MILVUS_DB_PATH)
print("Models and DB connection ready.")
然后将每页的PDF加载到一个文档对象中:
# Load PDF — one Document object per page
loader = PyPDFLoader(PDF_PATH)
documents = loader.load()
print(f"Loaded {len(documents)} pages.")
# Split into overlapping chunks
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP,
separators=["\n\n", "\n", ".", " ", ""]
)
chunks = text_splitter.split_documents(documents)
print(f"Created {len(chunks)} chunks.")
RecursiveCharacterTextSplitter 首先按段落分隔符进行分割,接着是换行符,然后是句子,这样就能尽可能将相关的句子保留在同一个块中。随后是两种向量类型的嵌入值:
texts = [doc.page_content for doc in chunks]
# Dense embeddings - one API call for the entire corpus
print("Generating dense embeddings...")
dense_embeddings = embedding_obj.embed_documents(texts)
print(f"Dense dimension: {len(dense_embeddings[0])}")
# Sparse embeddings - BM25 must be fit on YOUR corpus first
print("Fitting BM25 on corpus...")
analyzer = build_default_analyzer(language="en") # for this you will require nltk-stopwords
bm25_ef = BM25EmbeddingFunction(analyzer)
bm25_ef.fit(texts) # Builds vocabulary from your documents
sparse_embeddings = bm25_ef.encode_documents(texts)
print("Sparse embeddings generated.")
bm25_ef.fit(texts) 这一步至关重要。BM25 能从您自己的语料库中学习词汇及词频,也正是因此它才能发挥作用。若使用通用的 BM25 模型,关键词匹配的质量会明显下降。现在来看包含架构与索引的集合:
# Drop and recreate for a clean state
if COLLECTION_NAME in client.list_collections():
client.drop_collection(COLLECTION_NAME)
# Define schema
schema = client.create_schema()
schema.add_field("id", DataType.VARCHAR, is_primary=True, max_length=100)
schema.add_field("vector", DataType.FLOAT_VECTOR, dim=EMBEDDING_DIM)
schema.add_field("sparse_vector", DataType.SPARSE_FLOAT_VECTOR)
schema.add_field("text", DataType.VARCHAR, max_length=65535)
schema.add_field("page_number", DataType.INT64)
schema.add_field("source", DataType.VARCHAR, max_length=500)
schema.add_field("chunk_id", DataType.INT64)
# Create the collection
client.create_collection(collection_name=COLLECTION_NAME, schema=schema)
# Build indexes separately
index_params = client.prepare_index_params()
index_params.add_index(
field_name="vector",
index_type="FLAT", # Exact search - swap to HNSW for production
metric_type="COSINE"
)
index_params.add_index(
field_name="sparse_vector",
index_type="SPARSE_INVERTED_INDEX",
metric_type="IP" # Inner Product is the only valid metric for sparse
)
client.create_index(collection_name=COLLECTION_NAME, index_params=index_params)
print("Collection and indexes created.")
这里有两个索引选择需要考虑:
FLAT,用于密集向量。它能够进行精确的最近邻搜索,结果完全准确但处理大规模数据时速度较慢。在开发阶段以及数据集规模约为10万条记录以内时使用效果不错,而在生产环境中应选用HNSW替代它。SPARSE_INVERTED_INDEX,用于稀疏向量。这是Milvus中处理稀疏数据的唯一有效索引类型,其工作方式类似于传统的倒排索引,即搜索引擎通过单词对文档进行索引的方式。
最后,准备记录并将其插入:
def sparse_to_dict(s_emb) -> dict:
"""Convert a scipy sparse row into Milvus-compatible {index: value} dict."""
if sp.issparse(s_emb):
coo = s_emb.tocoo()
return {int(col): float(val) for col, val in zip(coo.col, coo.data)}
elif isinstance(s_emb, dict):
return s_emb
else:
return {int(i): float(v) for i, v in enumerate(s_emb) if v != 0.0}
# Build the records list
data = []
for idx, (chunk, d_emb) in enumerate(tqdm(zip(chunks, dense_embeddings), total=len(chunks))):
sparse_dict = sparse_to_dict(sparse_embeddings[idx])
if not sparse_dict:
print(f"Warning: empty sparse vector at chunk {idx}, skipping.")
continue
data.append({
"id": str(uuid.uuid4()),
"vector": d_emb,
"sparse_vector": sparse_dict,
"text": chunk.page_content,
"page_number": int(chunk.metadata.get("page", -1)),
"source": PDF_PATH,
"chunk_id": idx
})
# Insert into Milvus
res = client.insert(collection_name=COLLECTION_NAME, data=data)
print(f"Inserted {res['insert_count']} records.")
# Load into memory - required before any search operation
client.load_collection(COLLECTION_NAME)
print(f"Load state: {client.get_load_state(COLLECTION_NAME)}")
需要使用 sparse_to_dict 辅助函数,因为 BM25 会生成 scipy 稀疏矩阵,而 Milvus 需要的是普通的 Python {int: float} 字典。该函数能够处理二维稀疏行、一维数组以及已转换好的字典。一个常见的陷阱是:搜索前务必调用 client.load_collection()。在 Milvus 中,索引存储在磁盘上,会根据需求加载到内存中,因此如果不调用此函数,搜索将会出错或返回空结果。
以密集搜索作为基准
在数据已建立索引后,最简单的搜索方式是将查询嵌入其中,再通过余弦相似度找出距离最近的文档向量:
# ════════════════════════════════════════════════════════════
# Dense Vector Search
# ════════════════════════════════════════════════════════════
query = "What is the leave policy?"
# Step 1: Embed the query using the same model used at index time
query_dense_embedding = embedding_obj.embed_query(query)
# Step 2: Search
results = client.search(
collection_name=COLLECTION_NAME,
data=[query_dense_embedding],
anns_field="vector",
search_param={"metric_type": "COSINE"},
limit=TOP_K,
output_fields=["text", "page_number", "source"]
)
# Step 3: Display results
for idx, hit in enumerate(results[0], start=1):
entity = hit["entity"]
print(f"Rank {idx} | Cosine Score: {hit['distance']:.4f} | Page: {entity['page_number']}")
print(f" {entity['text'][:300]}\n")
在 COSINE 度量标准下,得分范围为 -1 到 1。通常而言,得分高于 0.7 表示语义匹配度很高,而最佳结果的平均得分多在 0.6 到 0.95 之间。对于那些答案位于相近语义空间内的自然语言问题,比如询问育儿假福利或行为准则的内容,密集搜索就已足够。但当查询包含嵌入模型很少见的具体代码、缩写或专有名词时,这种搜索方式就会显得力不从心。
结合 RRF 的混合搜索
混合搜索会同时执行密集搜索和稀疏的 BM25 搜索,然后通过 互反排名融合(RRF) 方法将两个排序列表合并。其原理很简单:每份文档的 RRF 得分是来自两个列表的 1 / (k + 排名) 值之和,其中 k=60 是用于平滑处理的常数,因此在这两个列表中排名都较高的文档最终会位于最顶端。
# ════════════════════════════════════════════════════════════
# Hybrid Search (Dense + Sparse)
# ════════════════════════════════════════════════════════════
query = "leave policy?"
# Dense query vector
query_dense = embedding_obj.embed_query(query)
# Sparse query vector - uses the same BM25 model fitted on the corpus
sparse_raw = bm25_ef.encode_queries([query])
sparse_dict = sparse_to_dict(sparse_raw[0])
print(f"Sparse query terms: {len(sparse_dict)}") # Should be > 0
# Build two separate ANN search requests
dense_req = AnnSearchRequest(
data=[query_dense],
anns_field="vector",
param={"metric_type": "COSINE"},
limit=TOP_K
)
sparse_req = AnnSearchRequest(
data=[sparse_dict],
anns_field="sparse_vector",
param={"metric_type": "IP"},
limit=TOP_K
)
# Execute hybrid search with RRF fusion
results = client.hybrid_search(
collection_name=COLLECTION_NAME,
reqs=[dense_req, sparse_req],
ranker=RRFRanker(k=60),
limit=TOP_K,
output_fields=["text", "page_number", "source"]
)
for idx, hit in enumerate(results[0], start=1):
entity = hit["entity"]
print(f"Rank {idx} | RRF Score: {hit['distance']:.4f} | Page: {entity['page_number']}")
print(f" {entity['text'][:300]}\n")
关于分数:hybrid_search 返回的 hit['distance'] 值是融合后的RRF排名分数,而非原始的余弦相似度值。数值越大越好,不过其范围并不限于0到1之间,因此不应与常规 search() 方法得到的余弦相似度分数直接对比。在实际应用中,混合搜索是推荐的默认方案。BM25几乎不会增加延迟,且能在包含大量关键词的查询中显著提升召回率。
即便采用混合检索方式,在某些特定情况下仍可能失败,有四种技术可用于解决最常见的故障模式。
元数据过滤
问题所在。当数据集涵盖多份文档、多个时期或部门时,向量搜索可能会从错误来源返回看似相关的内容,比如过时的政策版本、其他部门的手册,或是与问题范畴无关的页面。
工作原理。Milvus在向量搜索之前应用标量过滤表达式,只有匹配的行才会进入近似最近邻阶段。这不是事后过滤,而是通过过滤预先筛选候选集,从而在已过滤的子集中保持较高的召回率。
# ════════════════════════════════════════════════════════════
# METADATA FILTERING
# ════════════════════════════════════════════════════════════
def search_with_metadata_filter(
client, collection_name, embedding_obj, bm25_ef,
query: str,
page_range: tuple = None,
source_file: str = None,
top_k: int = 5
):
filter_parts = []
if page_range:
lo, hi = page_range
filter_parts.append(f"page_number >= {lo} && page_number <= {hi}")
if source_file:
filter_parts.append(f'source == "{source_file}"')
filter_expr = " && ".join(filter_parts) if filter_parts else None
print(f"\n[Metadata Filter] Query : '{query}'")
print(f"[Metadata Filter] Filter: {filter_expr or 'None (unfiltered)'}")
results = hybrid_search(
client, collection_name, embedding_obj, bm25_ef,
query_text=query,
top_k=top_k,
filters=filter_expr
)
return results
# ── Run ──────────────────────────────────────────────────────
meta_results = search_with_metadata_filter(
client, COLLECTION_NAME, embedding_obj, bm25_ef,
query = "What is the leave policy?",
page_range = (1, 30),
source_file= None,
top_k = 5
)
# ── Print Results ─────────────────────────────────────────────
print("\nMETADATA-FILTERED RESULTS")
print("=" * 55)
if not meta_results or not meta_results[0]:
print("No results returned.")
else:
for idx, hit in enumerate(meta_results[0], start=1):
entity = hit["entity"]
print(f"\nRank : {idx}")
print(f"Score : {hit['distance']:.4f}")
print(f"Page : {entity['page_number']}")
print(f"Text :\n{entity['text'][:400]}")
过滤表达式遵循以下语法:
'page_number >= 1 && page_number <= 30' # page range
'source == "hr_policy.pdf"' # exact source
'category in ["leave", "performance"]' # in a list
'source like "hr%"' # prefix match
任何多文档RAG系统都应具备元数据过滤功能。如果用户的工作范围有限,比如仅限于2024年手册或金融类文件,系统应自动根据该范围生成过滤条件。
查询重写
问题所在。用户通常会用较为随意的表述,例如“我有多少天假期?”,而文档中则使用正式用语,如“年假权益”或“休假天数累积政策”。密集搜索能在一定程度上弥补这一差距,而良好的重写方式则能显著提升检索效率与整体召回率。
工作原理。大型语言模型会为原始查询生成几种正式且包含大量关键词的重写版本,每个版本都会进行独立搜索,随后RRF机制会对搜索结果进行合并并去除重复内容。
# ════════════════════════════════════════════════════════════
# QUERY REWRITING
# ════════════════════════════════════════════════════════════
import re, json
REWRITE_PROMPT = """You are an expert at reformulating search queries to improve document retrieval.
Given a user query, produce {n} alternative search queries that:
- Use formal, document-style language
- Include relevant keywords and synonyms
- Cover different angles of the same question
User query: {query}
Respond ONLY with a JSON array of strings. Example:
["rewritten query 1", "rewritten query 2", "rewritten query 3"]"""
def rewrite_query(query: str, n: int = 3) -> list[str]:
prompt = REWRITE_PROMPT.format(query=query, n=n)
response = llm.invoke(prompt)
raw = re.sub(r"^```json|^```|```quot;, "", response.content.strip(), flags=re.MULTILINE).strip()
try:
variants = json.loads(raw)
return [query] + variants # always keep the original
except json.JSONDecodeError:
print("Warning: Could not parse rewrites, using original query only.")
return [query]
def search_with_query_rewriting(
client, collection_name, embedding_obj, bm25_ef,
query: str,
n_rewrites: int = 3,
top_k: int = 5
):
variants = rewrite_query(query, n=n_rewrites)
print(f"\n[Query Rewriting] Original : '{query}'")
for i, v in enumerate(variants[1:], 1):
print(f"[Query Rewriting] Variant {i} : '{v}'")
seen_ids = {}
rank_scores = {}
for variant in variants:
results = hybrid_search(
client, collection_name, embedding_obj, bm25_ef,
query_text=variant,
top_k=top_k
)
if not results or not results[0]:
continue
for rank, hit in enumerate(results[0], start=1):
hit_id = hit["id"]
rank_scores[hit_id] = rank_scores.get(hit_id, 0) + 1.0 / (60 + rank)
if hit_id not in seen_ids:
seen_ids[hit_id] = hit
merged = sorted(seen_ids.values(), key=lambda h: rank_scores[h["id"]], reverse=True)[:top_k]
return [merged]
# ── Run ──────────────────────────────────────────────────────
rewrite_results = search_with_query_rewriting(
client, COLLECTION_NAME, embedding_obj, bm25_ef,
query = "What is the leave policy?",
n_rewrites = 3,
top_k = 5
)
# ── Print Results ─────────────────────────────────────────────
print("\nQUERY-REWRITTEN RESULTS")
print("=" * 55)
if not rewrite_results or not rewrite_results[0]:
print("No results returned.")
else:
for idx, hit in enumerate(rewrite_results[0], start=1):
entity = hit["entity"]
print(f"\nRank : {idx}")
print(f"Score : {hit['distance']:.4f}")
print(f"Page : {entity['page_number']}")
print(f"Text :\n{entity['text'][:400]}")
一个转换示例:
Input: "how many days off do I get?"
Output variants:
1. "annual leave entitlement number of days employee handbook"
2. "vacation days accrual policy full-time employee"
3. "paid time off PTO allowance per calendar year"
每个变体都会扩大检索范围。在原始查询中排名前5之外的内容,在第三个变体中可能会排在第二位,而RRF能够呈现那些在各种查询中都持续出现的文档。
HyDE(假设性文档嵌入)
问题所在。简短的问题及其答案在嵌入空间中处于不同的位置。“有多少病假天数?”是一个问题,其嵌入向量会靠近其他问题;而规定员工每年可享受7天带薪病假的政策条款,其嵌入向量则会靠近其他政策声明。这种差异会降低检索的精确度。
工作原理。大语言模型不会直接嵌入问题,而是先撰写一份假设性的答案文档,再将该文档进行嵌入。即便其中的具体事实有误,这份编造的答案仍会出现在真实答案文档附近。系统仅使用其嵌入向量,而文本部分则被丢弃,因此这些虚构的细节永远不会传递给用户。
# ════════════════════════════════════════════════════════════
# HyDE (Hypothetical Document Embeddings)
# ════════════════════════════════════════════════════════════
HYDE_PROMPT = """You are a corporate policy document writer.
Write a 2-3 paragraph excerpt from an official HR policy or company document
that would DIRECTLY ANSWER the following question.
Write in formal document style. Do not mention the question itself.
Question: {query}
Document excerpt:"""
def generate_hypothetical_document(query: str) -> str:
response = llm.invoke(HYDE_PROMPT.format(query=query))
return response.content.strip()
def search_with_hyde(
client, collection_name, embedding_obj, bm25_ef,
query: str,
top_k: int = 5
):
hypothetical_doc = generate_hypothetical_document(query)
print(f"\n[HyDE] Query : '{query}'")
print(f"[HyDE] Hypothetical doc :\n {hypothetical_doc[:300]}...\n")
# Search using the hypothetical document's embedding
hyde_results = hybrid_search(
client, collection_name, embedding_obj, bm25_ef,
query_text=hypothetical_doc, # embed the answer, not the question
top_k=top_k
)
# Also search with the original query and merge both via RRF
original_results = hybrid_search(
client, collection_name, embedding_obj, bm25_ef,
query_text=query,
top_k=top_k
)
seen_ids = {}
rank_scores = {}
for result_set in [hyde_results, original_results]:
if not result_set or not result_set[0]:
continue
for rank, hit in enumerate(result_set[0], start=1):
hit_id = hit["id"]
rank_scores[hit_id] = rank_scores.get(hit_id, 0) + 1.0 / (60 + rank)
if hit_id not in seen_ids:
seen_ids[hit_id] = hit
merged = sorted(seen_ids.values(), key=lambda h: rank_scores[h["id"]], reverse=True)[:top_k]
return [merged]
# ── Run ──────────────────────────────────────────────────────
hyde_results = search_with_hyde(
client, COLLECTION_NAME, embedding_obj, bm25_ef,
query = "What is the leave policy?",
top_k = 5
)
# ── Print Results ─────────────────────────────────────────────
print("\nHyDE RESULTS")
print("=" * 55)
if not hyde_results or not hyde_results[0]:
print("No results returned.")
else:
for idx, hit in enumerate(hyde_results[0], start=1):
entity = hit["entity"]
print(f"\nRank : {idx}")
print(f"Score : {hit['distance']:.4f}")
print(f"Page : {entity['page_number']}")
print(f"Text :\n{entity['text'][:400]}")
HyDE在处理诸如薪资评估流程这类简短且模糊的查询、问题与文档之间的词汇差异,以及非常非正式的表达方式时效果最佳。对于那些包含大量精确关键词、混合搜索已能很好处理的查询,或是对LLM响应速度要求较高的场景,HyDE则显得过于复杂,因为每次查询都需要额外调用一次LLM。
查询分解
问题所在。复杂的多部分问题无法仅通过一段文本来解答。例如,同时询问休假政策以及绩效评估如何影响薪资,就涉及两个不同部分的内容。单一的向量搜索要么只关注其中一个主题而忽略另一个,要么对两个主题都只能给出一般性的结果。
工作原理。LLM会将问题拆分成多个独立的子问题,分别进行搜索,RRF负责合并这些结果,最终通过一次LLM调用获取所有相关的信息片段。
# ════════════════════════════════════════════════════════════
# QUERY DECOMPOSITION
# ════════════════════════════════════════════════════════════
DECOMPOSE_PROMPT = """You are an expert at breaking down complex questions for document retrieval.
Decompose the following question into 2-4 simple, self-contained sub-questions.
Each sub-question should target a single distinct piece of information.
Complex question: {query}
Respond ONLY with a JSON array of strings. Example:
["sub-question 1", "sub-question 2", "sub-question 3"]"""
def decompose_query(query: str) -> list[str]:
response = llm.invoke(DECOMPOSE_PROMPT.format(query=query))
raw = re.sub(r"^```json|^```|```quot;, "", response.content.strip(), flags=re.MULTILINE).strip()
try:
return json.loads(raw)
except json.JSONDecodeError:
print("Warning: Could not parse decomposition, using original query.")
return [query]
def search_with_decomposition(
client, collection_name, embedding_obj, bm25_ef,
query: str,
top_k: int = 5
):
sub_questions = decompose_query(query)
print(f"\n[Decomposition] Original query : '{query}'")
for i, sq in enumerate(sub_questions, 1):
print(f"[Decomposition] Sub-question {i} : '{sq}'")
per_subquery_results = {}
seen_ids = {}
rank_scores = {}
for sq in sub_questions:
results = hybrid_search(
client, collection_name, embedding_obj, bm25_ef,
query_text=sq,
top_k=top_k
)
per_subquery_results[sq] = results
if not results or not results[0]:
continue
for rank, hit in enumerate(results[0], start=1):
hit_id = hit["id"]
rank_scores[hit_id] = rank_scores.get(hit_id, 0) + 1.0 / (60 + rank)
if hit_id not in seen_ids:
seen_ids[hit_id] = hit
merged = sorted(seen_ids.values(), key=lambda h: rank_scores[h["id"]], reverse=True)[:top_k]
# Per sub-question breakdown
print("\n── Per Sub-question Results ──")
for sq, res in per_subquery_results.items():
print(f"\n SUB-QUERY: '{sq[:60]}'")
if res and res[0]:
for i, hit in enumerate(res[0], start=1):
print(f" {i}. Page {hit['entity']['page_number']} | Score {hit['distance']:.4f} | {hit['entity']['text'][:150]}")
return {"per_subquery": per_subquery_results, "merged": [merged]}
# ── Run ──────────────────────────────────────────────────────
decomp_results = search_with_decomposition(
client, COLLECTION_NAME, embedding_obj, bm25_ef,
query = "What is the leave policy and how does it affect salary deductions?",
top_k = 5
)
# ── Print Merged Results ──────────────────────────────────────
print("\nDECOMPOSED — MERGED FINAL RESULTS")
print("=" * 55)
merged_hits = decomp_results["merged"]
if not merged_hits or not merged_hits[0]:
print("No results returned.")
else:
for idx, hit in enumerate(merged_hits[0], start=1):
entity = hit["entity"]
print(f"\nRank : {idx}")
print(f"Score : {hit['distance']:.4f}")
print(f"Page : {entity['page_number']}")
print(f"Text :\n{entity['text'][:400]}")
一个分解示例:
Input: "What is the leave policy and how does performance review affect salary?"
Sub-questions:
1. "What is the annual leave policy?"
2. "How many sick days are employees entitled to?"
3. "How does performance review affect salary?"
4. "What is the performance review schedule?"
每个子问题都会独立检索所需的段落,这样合并后的结果就能覆盖原问题的所有部分,这是单一检索方式无法做到的。
交叉编码器重排序
目前的所有检索都是通过双编码器实现的:查询和文档被独立嵌入,相似度通过点积或余弦距离来计算。这种方法虽然快速,但并非评估相关性的最佳方式。交叉编码器则将查询和一条文档作为一对一起处理,通过观察它们之间的完整交互来评估匹配度,因此准确度要高得多。不过这种方法的速度较慢,处理数千份文档时效率很低,但在20到50份候选文档的小规模场景下则表现迅速。正因如此,通常采用先通过双编码器进行广泛检索,再通过交叉编码器对结果进行精细重排序的策略。
# ============================================================
# RERANKING WITH CROSS-ENCODER
# ============================================================
from sentence_transformers import CrossEncoder
# Huggingface: cross-encoder/ms-marco-MiniLM-L12-v2
cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L12-v2")
query = "What is the leave policy?"
RETRIEVAL_K = 20 # fetch more than you need
FINAL_K = 5 # rerank down to this
# Step 1: Broad retrieval - fetch 20 candidates
query_dense = embedding_obj.embed_query(query)
results = client.search(
collection_name=COLLECTION_NAME,
data=[query_dense],
anns_field="vector",
search_param={"metric_type": "COSINE"},
limit=RETRIEVAL_K,
output_fields=["text", "page_number", "source"]
)
hits = results[0]
print(f"Retrieved {len(hits)} candidates for reranking.")
# Step 2: Score each (query, chunk) pair with the cross-encoder
pairs = [[query, hit["entity"]["text"]] for hit in hits]
rerank_scores = cross_encoder.predict(pairs)
# Step 3: Sort by cross-encoder score
for hit, score in zip(hits, rerank_scores):
hit["rerank_score"] = float(score)
reranked = sorted(hits, key=lambda x: x["rerank_score"], reverse=True)[:FINAL_K]
# Step 4: Display
for idx, hit in enumerate(reranked, start=1):
entity = hit["entity"]
print(f"Rank {idx} | Rerank: {hit['rerank_score']:.4f} | Vector: {hit['distance']:.4f}")
print(f" Page {entity['page_number']}: {entity['text'][:300]}\n")
交叉编码器的得分是未标准化的对数几率值,而非概率值。数值越高表示相关性越强,因此8.0远高于2.0,但不同模型之间的绝对数值尺度存在差异。所使用的模型ms-marco-MiniLM-L12-v2是一种强大且轻量级的选择,它在MS MARCO段落排序数据集上经过训练,能够很好地处理针对文档段落的问题式查询。
各技术的对比
一种合理的生产环境默认配置为:
- 将混合搜索作为始终启用的基础检索方法。
- 在能够确定用户查询范围时,使用元数据过滤。
- 在将前20个候选结果传递给大语言模型之前,进行交叉编码器重排序。
其发展流程如下:
- 密集搜索是起点:速度快、无需依赖关键词,且对同义词处理良好。
- 混合搜索是升级版本,因为引入BM25稀疏向量能在几乎不增加复杂度的情况下有效提升召回率。
- 过滤、重写、HyDE和分解分别针对不同的故障模式。不要盲目全部使用,而应根据实际文档及用户查询模式来选择。
- 交叉编码器重排序几乎总是值得添加的。它在本地运行,对少量候选项的延迟影响很小,且能可靠地提升最终结果的质量。
归根结底,检索质量决定了答案的质量。大型语言模型的表现仅取决于其所获得的上下文质量,而上述各种技术都是为了让这些上下文更具相关性、更完整且更可靠。所有操作都在 Milvus Lite 上本地完成,仅依赖 OpenAI API 作为嵌入生成功能的外部接口。