稀疏向量检索流程

import jieba
from rank_bm25 import BM25Okapi
from langchain_text_splitters import RecursiveCharacterTextSplitter
from collections import Counter

# ──────────────────────────────
# 1. 文档预处理 + 分词
# ──────────────────────────────
docs = [
    "苹果手机最新功能包括AI摄影和长续航。发布会上提到了端侧AI与影像算法升级。",
    "香蕉是一种热带水果，富含钾元素，适合作为日常补充能量的食物。",
    "iPhone 16 Pro Max评测：屏幕更亮，影像更强，续航也有提升。",
]

# 中文分词函数（jieba 默认精确模式）
def tokenize_cn(text):
    return [t.strip() for t in jieba.cut(text) if t.strip()]

# ──────────────────────────────
# 辅助函数：生成稀疏向量
# ──────────────────────────────
def bm25_doc_sparse_vector(bm25: BM25Okapi, doc_index: int) -> dict[str, float]:
    """为指定 chunk 生成 BM25 稀疏向量（每个 token → 权重）"""
    freqs = bm25.doc_freqs[doc_index]
    dl = bm25.doc_len[doc_index]
    denom_const = bm25.k1 * (1.0 - bm25.b + bm25.b * dl / bm25.avgdl)  # 长度归一化项
    vec: dict[str, float] = {}
    for token, tf in freqs.items():
        idf = bm25.idf.get(token)
        if idf is None:
            continue
        score = idf * (tf * (bm25.k1 + 1.0)) / (tf + denom_const)  # 经典 BM25 TF 饱和公式
        if score != 0.0:
            vec[token] = float(score)
    return vec

def bm25_query_sparse_vector(bm25: BM25Okapi, query_tokens: list[str]) -> dict[str, float]:
    """为查询生成稀疏向量（每个 token → IDF × query_tf）"""
    qtf = Counter(query_tokens)
    vec: dict[str, float] = {}
    for token, tf in qtf.items():
        idf = bm25.idf.get(token)
        if idf is None:
            continue
        vec[token] = float(idf * tf)
    return vec

def format_sparse_vector(vec: dict[str, float], vocab_index: dict[str, int], top_k: int = 20) -> list[tuple[int, float, str]]:
    """将稀疏向量格式化为 (vocab_id, weight, token) 三元组，按权重降序"""
    items = sorted(vec.items(), key=lambda x: x[1], reverse=True)[:top_k]
    return [(vocab_index[t], w, t) for t, w in items]

# ──────────────────────────────
# 分块（Chunking）
# ──────────────────────────────
splitter = RecursiveCharacterTextSplitter(
    chunk_size=60,
    chunk_overlap=12,
    separators=["\n\n", "\n", "。", "！", "？", "，", " ", ""],
)
chunks = []
chunk_sources = []
for i, doc in enumerate(docs):
    for chunk in splitter.split_text(doc):
        if chunk.strip():
            chunks.append(chunk)
            chunk_sources.append(i)

# ──────────────────────────────
# 2. 构建 BM25 索引
# ──────────────────────────────
tokenized_chunks = [tokenize_cn(chunk) for chunk in chunks]
print(tokenized_chunks)
bm25 = BM25Okapi(tokenized_chunks)
vocab_index = {token: i for i, token in enumerate(sorted(bm25.idf.keys()))}

# ──────────────────────────────
# 3. 查询匹配
# ──────────────────────────────
query = "苹果手机最新功能"
tokenized_query = tokenize_cn(query)
scores = bm25.get_scores(tokenized_query)
top_chunks = bm25.get_top_n(tokenized_query, chunks, n=3)

print("Top chunks:", top_chunks)
print("Query sparse vector:", format_sparse_vector(bm25_query_sparse_vector(bm25, tokenized_query), vocab_index, top_k=50))
for i, chunk in enumerate(chunks):
    vec = bm25_doc_sparse_vector(bm25, i)
    print(f"Chunk[{i}] source_doc={chunk_sources[i]} score={float(scores[i]):.6f} text={chunk}")
    print("Sparse vector:", format_sparse_vector(vec, vocab_index, top_k=20))

{'苹果': 0.5108, '手机': 0.5108, 'AI': 0.5108, '续航': 0.0989, '影像': 0.0989, '。': 0.0989, ...}

[
  ['苹果', '手机', '最新', '功能', '包括', 'AI', '摄影', '和', '长', '续航', '。', '发布会', '上', '提到', '了', '端侧', 'AI', '与', '影像', '算法', '升级', '。'],
  ['香蕉', '是', '一种', '热带', '水果', '，', '富含', '钾', '元素', '，', '适合', '作为', '日常', '补充', '能量', '的', '食物', '。'],
  ['iPhone', '16', 'Pro', 'Max', '评测', '：', '屏幕', '更亮', '，', '影像', '更强', '，', '续航', '也', '有', '提升', '。']
]

Query: [(38, 0.5108, '苹果'), (21, 0.5108, '手机'), (29, 0.5108, '最新'), (13, 0.5108, '功能')]

Chunk[0]: [(AI, 0.6945), (苹果, 0.4769), (手机, 0.4769), (最新, 0.4769), ...]