Avatar

Qianqiu

Master in PKU Research Direction: VLM, RLHF, MLsys Hobbies: Game, Web novel, Anime

  1. WeChat
  1. Home
  2. Diary
  3. Research
  4. Entertain
  5. Search
  6. Archives
  7. About
    1. Dark Mode Light Mode

Table of contents

    1. embeddings.py
    2. vector_store.py
      1. _create_index():按类型创建 FAISS 索引
      2. build_from_documents():建库全流程
      3. search():在线检索流程
      4. add_documents():增量添加
      5. 总结
    3. sparse_store.py
    4. colbert_store.py
    5. retriever.py
      1. 三种检索如何工作
    6. llm.py
    7. rag_chain.py
转码日记

学习day31:熬夜睡醒就18:00 晚上学了一点

Feb 20, 2026

lc:763

重新默写复习了MHA MQA GQA

接着学习RAG的项目

embeddings.py

用 FlagEmbedding 的 BGEM3FlagModel 加载 BGE-M3。

  • get_embeddings() 工厂函数:按字符串选择用哪个模型。
  • __main__ 里演示:编码文本、打印向量维度、做余弦相似度测试。

BGEM3Embeddings:核心封装类(LangChain Embeddings 接口) self._load_model() 优先用 FlagEmbedding,失败则 fallback _load_fallback_model():用 sentence-transformers 兜底

embed_documents:批量把文本变成向量(建库时用)

encode_multi / encode_query_multi:同时拿 dense + sparse + ColBERT(仅 FlagEmbedding):要用 BGE-M3 的“三合一能力”时(dense/sparse/colbert)才用它。

get_dimension:动态测维度

get_embeddings 工厂函数:统一入口 - 好处:外部代码不用关心具体类,传一个字符串就行。 常用在配置驱动的项目里:比如 dev 用 simple,prod 用 bge-m3。

vector_store.py

把“文档 → 向量 → FAISS 索引 → 检索/保存/加载/增量添加”封装成一个类

一个类:FAISSVectorStore 核心能力 4 件套:

  1. build_from_documents:从一堆 Document 构建索引(一次性建库)
  2. search:给定 query 做 TopK 相似度检索(在线检索)
  3. save/load:把索引和文档元数据落盘 & 再读回来(持久化)
  4. add_documents:增量添加新文档(更新库)

_create_index():按类型创建 FAISS 索引

FlatIP:暴力遍历计算内积,精确,简单可靠 IVFFlat(近似检索):

  • 先把所有向量聚成 nlist 个簇(桶)
  • 查询时先找最可能的 nprobe 个桶
  • 只在这些桶里做精确比较 → 速度大幅提升,精度会有点损失 nprobe 越大:
  • 搜得更全(精度/召回更高)
  • 但速度更慢 经验式的 nlist 选择
  • num_vectors // 40:大概每个桶 40 个向量(粗略经验)
  • 同时限制不超过配置的 FAISS_NLIST HNSW(图近似检索):用图结构做近似最近邻,通常 召回高、速度也快.内存占用更大

build_from_documents():建库全流程

提取文本 向量化 更新实际维度 归一化(归一化后,内积等价余弦相似度,所以后面用 IndexFlatIP / METRIC_INNER_PRODUCT) 创建索引 训练(仅 IVF 需要) 把向量加进去

search():在线检索流程

query 向量化 + 归一化 FAISS 搜索 scores, indices = self.index.search(query_vector, top_k) 过滤与组装结果

1
2
3
4
for score, idx in zip(scores[0], indices[0]):  
if idx == -1: continue  
if score < score_threshold: continue  
results.append((self.documents[idx], float(score)))

add_documents():增量添加

流程基本是 build 的缩小版,仅IVFFLAT有细节处理上的问题。

  • IVFFlat 在添加前必须 train 过。如果你没 build 直接 add,会报错。
  • 如果“先 build 很少数据训练 IVF,再大量增量 add”,效果可能不好(聚类中心不适配新分布)。实践中可能需要重建/再训练。

总结

必须对 doc/query 都做 L2 normalize,否则内积不是余弦,阈值/分数含义会乱。 documents 顺序必须和向量添加顺序一致,否则 self.documents[idx] 会错配。 IVFFlat 必须 train 后才能 add,并且训练数据量太少会影响质量。

sparse_store.py

稀疏向量表示,比传统 BM25 更准确:

  • 模型为每个 token 学习了重要性权重

  • 可以捕获语义同义词(BM25 做不到)

  • 格式:{token_id_str: weight, …}

检索原理:对 query 和 document 的稀疏向量做内积(只遍历 query 中非零 token)

细节略过

colbert_store.py

ColBERT (Contextualized Late Interaction over BERT) 原理:

  • 每个文档不再用单个向量表示,而是保留每个 token 的向量
  • 查询时对每个 query token 找文档中最相似的 token,求和:   score = Σ_i max_j cos(q_i, d_j)
  • 精度高于单向量检索,但计算开销大
  • 因此通常只对候选集做重打分,不做全库检索

retriever.py

拿到候选文档后(可选)再做 Cross-Encoder 重排序,最终返回 Top-K。

BM25Retriever 用 jieba 对中文分词,把每个文档切成词序列。

  • 用 BM25Okapi 建索引。
  • search(query):同样分词后 BM25 打分,取分数最高的 top_k 文档

Reranker 输入是初检结果,重新按 rerank 分排序,返回 top_k。

HybridRetriever(总控) 记录各种权重与模式参数

三种检索如何工作

dense_only FAISS dense 初检用 top_k * 2:给 reranker 留更大候选池,提高最终质量。

hybrid_bm25

  • Dense:FAISS 搜 top_k*2
  • BM25:bm25 搜 top_k*2
  • _fuse_results:两路 min-max 归一化到 [0,1] 后加权相加
  • final = dense_weight * norm_dense + bm25_weight * norm_bm25
  • 过滤阈值 + (可选)rerank → top_k

multi_route query_multi = embeddings.encode_query_multi(query) 得到 dense_vecs / lexical_weights / colbert_vecs

  • 合并成 candidate_set[doc_idx] = {doc, dense_score, sparse_score} ColBERT 对候选集重打分(token 级 MaxSim)
  • 分别对 dense/sparse/colbert 的分数做 min-max 归一化
  • 然后三路加权相加:
    fused = w_dense*dense_norm + w_sparse*sparse_norm + w_colbert*colbert_norm
  • 排序得到融合结果
  • 回到 retrieve():阈值过滤 + (可选)rerank → top_k 从 candidate_set 收集三路分数列表 _normalize_score_list 做 min-max 按候选顺序逐个算 fused_score,最后排序

llm.py

把本地 Ollama 部署的 Qwen(如 qwen2.5:7b)通过 LangChain 的 ChatOllama 包装成一个统一的“聊天模型接口”,并提供一个连通性测试和一个可运行示例

rag_chain.py

RAGChain 类:初始化 __init__ 做了 4 件事

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
self.vector_store = vector_store
self.llm = llm or get_llm()
self.top_k = top_k
self.retriever = create_retriever(...)
self.chat_history: list = []
self.prompt = ChatPromptTemplate.from_messages([  
("system", SYSTEM_PROMPT),  
MessagesPlaceholder(variable_name="chat_history", optional=True),  
("human", QA_PROMPT_TEMPLATE),  
])

保存核心依赖 创建检索器 初始化对话历史 构建 Prompt 模板

query():非流式的完整 RAG 问答

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
retrieved = self.retriever.retrieve(query=question, top_k=top_k)
context = self._format_context(retrieved)
messages = self.prompt.format_messages(  
context=context,  
question=question,  
chat_history=self.chat_history[-6:], # 最近3轮  
)
response = self.llm.invoke(messages)  
answer = response.content
self.chat_history.append(HumanMessage(content=question))  
self.chat_history.append(AIMessage(content=answer))
sources = self._extract_sources(retrieved)

stream():流式输出版本 整体和 query() 一样,差别在 LLM 调用

_format_context():上下文组装规则 它把每条文档变成“可被引用”的块

_extract_sources():做“参考来源列表” - 去重:

  • 同 source 只保留一次
  • 给一个 preview:让用户知道引用了什么内容
  • score 取 round(4)

create_rag_chain():一键拉起整套系统

interactive_mode():命令行聊天入口

Related content

学习day46:尘埃落定

学习day45:休养生息,也许接近这次的转码日记尾声了,期待下次再出发。

学习day44:接到第一个offer,晚上又面试,间歇性学习八股和力扣

学习day43:午晚各一面

学习day42:一面水过去

© 2025 - 2026 Qianqiu
Built with Hugo
Theme Stack designed by Jimmy