今天探讨的便是RAG开发关于全文搜索，以及将全文检索和向量检索整合在一起的混合搜索（Hybrid Search）。

全文搜索（BM25） 代码实现

pip install llama-index-retrievers-bm25

为了便于在DuckDB里重现全文索引，这里使用DuckDB来存储

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader,Settings
from llama_index.llms.ollama import Ollama
from llama_index.core.node_parser import SentenceSplitter
from llama_index.retrievers.bm25 import BM25Retriever


# 这里仍然使用Ollama，你可以根据自己的情况，选择合适的LLM和embedding model
Settings.llm = Ollama(model="wizardlm2:7b-q5_K_M", request_timeout=300.0)

from llama_index.embeddings.ollama import OllamaEmbedding
Settings.embed_model = OllamaEmbedding(model_name="snowflake-arctic-embed:latest")
embed_dim=len(Settings.embed_model.get_query_embedding('hello'))

# 索引
documents = SimpleDirectoryReader("./data/paul_graham").load_data()
splitter = SentenceSplitter(chunk_size=1024) # 为了bm25需要，可以适当改小该值
nodes = splitter.get_nodes_from_documents(documents)

vector_store = DuckDBVectorStore(embed_dim=embed_dim,
                                 database_name="paul.duck",
                                 persist_dir="duckdb")
storage_context = StorageContext.from_defaults(vector_store=vector_store)
storage_context.docstore.add_documents(nodes)

# vector retriever
vector_retriever = index.as_retriever(similarity_top_k=5)

# BM25Retriever
bm25_retriever = BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=5)

query = "What happened at Viaweb and Interleaf?"

from llama_index.core.response.notebook_utils import display_source_node

# will retrieve context from specific companies
nodes_bm25 = bm25_retriever.retrieve(query)
for node in nodes_bm25:
    display_source_node(node)

使用DuckDB的全文索引来重现BM25Retriever

import duckdb
conn = duckdb.connect('duckdb/paul.duck')
conn.sql('select * from documents limit 5').df()

conn.sql('install fts') # 第一次需要安装全文索引插件
conn.sql('load fts') # 加载全文索引插件

# 创建全文索引

conn.sql("PRAGMA create_fts_index('documents', 'node_id', 'text');")
# 重现 bm25_retriever.retrieve(query)
conn.sql(f"""
SELECT
    fts_main_documents.match_bm25(node_id, '{query}') AS score,
    node_id, text
FROM documents
WHERE score IS NOT NULL
ORDER BY score DESC
LIMIT 5
""").df()

对比bm25_retriever.retrieve(query)的实现， 结果略有不同， 因为使用的bm25引擎不同， bm25_retriever 使用的是 rank_bm25^[1]

关于DuckDB 全文索引插件的使用，详见，https://duckdb.org/docs/extensions/full_text_search

混合搜索（Hybrid Search）

接下来重现全文检索与向量检索选哪个？ 提到的混合搜索

相对得分融合

relativeScoreFusion 通过分别归一化向量搜索和关键词搜索输出的指标来确定每个对象的得分。最高值设为1，最低值设为0，其他值根据这个尺度介于两者之间。因此，总得分是通过归一化的向量相似度和归一化的BM25得分的加权和来计算的。

from llama_index.core.retrievers import QueryFusionRetriever

retriever_relative_score = QueryFusionRetriever(
    [vector_retriever, bm25_retriever],
    retriever_weights=[0.6, 0.4],
    similarity_top_k=10,
    num_queries=1,  # set this to 1 to disable query generation
    mode="relative_score",
    use_async=True,
    verbose=True,
)
nodes_with_scores = retriever.retrieve(
    "What happened at Interleafe and Viaweb?"
)
for node in nodes_with_scores:
    print(f"Score: {node.score:.2f} - {node.text[:100]}...\n-----")

from llama_index.core.query_engine import RetrieverQueryEngine

query_engine = RetrieverQueryEngine.from_args(retriever_relative_score)
response = query_engine.query("What happened at Interleafe and Viaweb?")
display_response(response)

分布的得分融合

相对得分融合（Relative Score Fusion）的一个变体，基于分布的得分融合（Distribution-Based Score Fusion）以稍微不同的方式对得分进行缩放——基于每个结果集的得分的平均值和标准差。

from llama_index.core.retrievers import QueryFusionRetriever

retriever = QueryFusionRetriever(
    [vector_retriever, bm25_retriever],
    retriever_weights=[0.6, 0.4],
    similarity_top_k=10,
    num_queries=1,  # set this to 1 to disable query generation
    mode="dist_based_score",
    use_async=True,
    verbose=True,
)

nodes_with_scores = retriever.retrieve(
    "What happened at Interleafe and Viaweb?"
)

for node in nodes_with_scores:
    print(f"Score: {node.score:.2f} - {node.text[:100]}...\n-----")

from llama_index.core.query_engine import RetrieverQueryEngine

query_engine = RetrieverQueryEngine.from_args(retriever_relative_score)
response = query_engine.query("What happened at Interleafe and Viaweb?")
display_response(response)

结论

到现在，我们不仅仅可以进行向量搜索，可以进行全文搜索，还可以使用两种方法进行混合搜索了。

引用链接

[1] rank_bm25: https://github.com/dorianbrown/rank_bm25