如何用Haystack与Faiss匹配Top k句子并返回对应文档？

我有1000份纯文本文档和一个用户查询，想要用Python库Haystack和Faiss检索最相关的Top k文档。具体需求是：先识别出与用户查询最匹配的Top k句子，再返回包含这些句子的文档。现有代码能识别与查询最匹配的Top k文档，请问如何修改代码实现上述需求？是否可以调整Faiss的索引策略？

现有代码：

# Note: Most of the code is from https://haystack.deepset.ai/tutorials/07_rag_generator

import logging

logging.basicConfig(format="%(levelname)s - %(name)s -  %(message)s", level=logging.WARNING)
logging.getLogger("haystack").setLevel(logging.INFO)

import pandas as pd
from haystack.utils import fetch_archive_from_http

# Download sample
doc_dir = "data/tutorial7/"
s3_url = "https://s3.eu-central-1.amazonaws.com/deepset.ai-farm-qa/datasets/small_generator_dataset.csv.zip"
fetch_archive_from_http(url=s3_url, output_dir=doc_dir)

# Create dataframe with columns "title" and "text"
#df = pd.read_csv(f"{doc_dir}/small_generator_dataset.csv", sep=",")
df = pd.read_csv(f"{doc_dir}/small_generator_dataset.csv", sep=",",nrows=10)
# Minimal cleaning
df.fillna(value="", inplace=True)

print(df.head())

from haystack import Document

# Use data to initialize Document objects
titles = list(df["title"].values)
texts = list(df["text"].values)
documents = []
for title, text in zip(titles, texts):
    documents.append(Document(content=text, meta={"name": title or ""}))

from haystack.document_stores import FAISSDocumentStore
document_store = FAISSDocumentStore(faiss_index_factory_str="Flat", return_embedding=True)

from haystack.nodes import RAGenerator, DensePassageRetriever

retriever = DensePassageRetriever(
    document_store=document_store,
    query_embedding_model="facebook/dpr-question_encoder-single-nq-base",
    passage_embedding_model="facebook/dpr-ctx_encoder-single-nq-base",
    use_gpu=True,
    embed_title=True,
)

# Delete existing documents in documents store
document_store.delete_documents()

# Write documents to document store
document_store.write_documents(documents)

# Add documents embeddings to index
document_store.update_embeddings(retriever=retriever)

from haystack.pipelines import GenerativeQAPipeline
from haystack import Pipeline
pipeline = Pipeline()
pipeline.add_node(component=retriever, name='Retriever', inputs=['Query'])

from haystack.utils import print_answers

QUESTIONS = [
    "who got the first nobel prize in physics",
    "when is the next deadpool movie being released",
]

for question in QUESTIONS:
    res = pipeline.run(query=question, params={"Retriever": {"top_k": 5}})
    print(res)
    #print_answers(res, details="all")

原代码运行指令：

conda create -y --name haystacktest python==3.9
conda activate haystacktest
pip install --upgrade pip
pip install farm-haystack
conda install pytorch -c pytorch
pip install sentence_transformers
pip install farm-haystack[colab,faiss]==1.17.2

核心思路 #

要实现需求，关键是将原始文档拆分为句子级别的小单元，先检索匹配度最高的Top k句子，再通过元数据关联回对应的原始文档。同时可根据数据规模调整Faiss索引策略，平衡检索速度与精度。

具体实现步骤 #

1. 句子级文档拆分：用NLP工具将每个原始文档拆分为独立句子，每个句子作为单独的Document对象，并在元数据中记录所属原始文档的标识（如标题、索引）。
2. 检索逻辑调整：先检索Top k句子，再从这些句子的元数据中提取原始文档信息，去重后返回结果。
3. Faiss索引优化：当拆分后的句子数量较大（如十万级以上），将默认的Flat精确索引替换为近似索引（如IVF_Flat、HNSW）提升检索效率。

修改后的完整代码 #

# 基于haystack官方教程修改
import logging
import nltk
nltk.download('punkt')
from nltk.tokenize import sent_tokenize

logging.basicConfig(format="%(levelname)s - %(name)s -  %(message)s", level=logging.WARNING)
logging.getLogger("haystack").setLevel(logging.INFO)

import pandas as pd
from haystack.utils import fetch_archive_from_http

# 下载示例数据
doc_dir = "data/tutorial7/"
s3_url = "https://s3.eu-central-1.amazonaws.com/deepset.ai-farm-qa/datasets/small_generator_dataset.csv.zip"
fetch_archive_from_http(url=s3_url, output_dir=doc_dir)

# 读取数据
df = pd.read_csv(f"{doc_dir}/small_generator_dataset.csv", sep=",", nrows=10)
df.fillna(value="", inplace=True)

from haystack import Document

# 将原始文档拆分为句子级Document，记录原文档关联信息
sentence_docs = []
for idx, row in df.iterrows():
    title = row["title"]
    text = row["text"]
    # 拆分句子
    sentences = sent_tokenize(text)
    for sent_idx, sentence in enumerate(sentences):
        sentence_docs.append(
            Document(
                content=sentence,
                meta={
                    "original_title": title,
                    "original_doc_index": idx,
                    "sentence_index": sent_idx
                }
            )
        )

from haystack.document_stores import FAISSDocumentStore
# 调整Faiss索引策略，IVF_Flat适合大规模数据，平衡速度与精度
document_store = FAISSDocumentStore(
    faiss_index_factory_str="IVF1024,Flat",
    return_embedding=True,
    similarity="dot_product"  # DPR模型推荐使用点积相似度
)

from haystack.nodes import DensePassageRetriever

retriever = DensePassageRetriever(
    document_store=document_store,
    query_embedding_model="facebook/dpr-question_encoder-single-nq-base",
    passage_embedding_model="facebook/dpr-ctx_encoder-single-nq-base",
    use_gpu=True,
    embed_title=False  # 句子级无需嵌入原标题，若标题有价值可改为True
)

# 清空文档库并写入句子级文档
document_store.delete_documents()
document_store.write_documents(sentence_docs)
document_store.update_embeddings(retriever=retriever)

from haystack import Pipeline

pipeline = Pipeline()
pipeline.add_node(component=retriever, name='Retriever', inputs=['Query'])

# 自定义函数：从Top k句子映射回原始文档并去重
def get_top_docs_from_top_sentences(query, top_k_sentences=5):
    # 检索Top k匹配句子
    res = pipeline.run(query=query, params={"Retriever": {"top_k": top_k_sentences}})
    retrieved_sentences = res["documents"]
    
    # 收集原始文档信息并去重
    original_docs = {}
    for sent_doc in retrieved_sentences:
        doc_idx = sent_doc.meta["original_doc_index"]
        if doc_idx not in original_docs:
            original_docs[doc_idx] = {
                "title": sent_doc.meta["original_title"],
                "text": df.iloc[doc_idx]["text"],
                "matched_sentences": [sent_doc.content]
            }
        else:
            original_docs[doc_idx]["matched_sentences"].append(sent_doc.content)
    
    # 转换为列表，保留检索优先级
    result_docs = list(original_docs.values())
    return result_docs

# 测试查询
QUESTIONS = [
    "who got the first nobel prize in physics",
    "when is the next deadpool movie being released",
]

for question in QUESTIONS:
    print(f"\n=== 查询: {question} ===")
    top_docs = get_top_docs_from_top_sentences(question, top_k_sentences=5)
    for i, doc in enumerate(top_docs):
        print(f"\nTop {i+1} 文档:")
        print(f"标题: {doc['title']}")
        print(f"匹配句子: {', '.join(doc['matched_sentences'])}")
        print(f"完整内容预览: {doc['text'][:200]}...")

调整后的运行指令 #

conda create -y --name haystacktest python==3.9
conda activate haystacktest
pip install --upgrade pip
pip install farm-haystack
conda install pytorch -c pytorch
pip install sentence_transformers nltk
pip install farm-haystack[colab,faiss]==1.17.2

关键说明 #

• 句子拆分工具：示例用NLTK的sent_tokenize，也可替换为spaCy等工具提升多语言拆分精度。
• Faiss索引选择：
  • 数据量较小（<10万句子）：用Flat索引，保证精确匹配
  • 数据量中等（10万-100万句子）：用IVF_Flat，通过聚类减少搜索范围
  • 数据量极大（>100万句子）：用HNSW，检索速度最快，内存占用略高
• 去重逻辑：同一个原始文档可能有多个句子被检索到，通过原文档索引去重并合并匹配句子，避免返回重复文档。

内容的提问来源于stack exchange，提问作者Franck Dernoncourt