使用 BERT 的 NSP 实现语义感知切片 —— 提升 RAG 系统的检索质量

在构建 Retrieval-Augmented Generation(RAG)系统时,文档的切片方式至关重要。我们需要将长文本切分成合适的段落(chunks),然后存入向量数据库进行召回。如果切得太粗,会丢失上下文细节;切得太细,语义就会支离破碎。

传统切片方法(如固定 token 窗口、按句子数切)不考虑句子间的语义关系,很容易导致把两个紧密相关的句子切开,影响后续检索和生成质量。

为了解决这个问题,我们可以借助 BERT 模型中的 Next Sentence Prediction(NSP)机制 ,实现语义感知的切片

什么是 NSP?

NSP 是 BERT 在预训练阶段的两个任务之一,其目标是判断两个句子是否在原始文本中是相邻的。具体来说,给定句子 A 和 B,BERT 会输出一个判断:

"B 是不是 A 的下一句?"

我们可以反过来利用这个能力 ,判断两个句子之间是否具有语义连续性,从而找出应该切片的位置

如何用 NSP 做文本切片?

步骤概览:

  1. 将文本分句(sentence tokenization)

  2. 对每对相邻句子使用 BERT NSP 判断是否连贯

  3. 如果不连贯 → 作为一个切片边界

  4. 合并成语义一致的 chunk,用于向量化和检索

Python 实现示例

python 复制代码
from transformers import BertTokenizer, BertForNextSentencePrediction
import torch
from nltk.tokenize import sent_tokenize
import nltk

# 下载 NLTK 分句模型
nltk.download('punkt')

# 初始化 BERT NSP 模型
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForNextSentencePrediction.from_pretrained('bert-base-uncased')
model.eval()

# 判断两个句子是否是"语义上连贯"
def is_next_sentence(sent1, sent2, threshold=0.5):
    inputs = tokenizer.encode_plus(sent1, sent2, return_tensors='pt')
    with torch.no_grad():
        logits = model(**inputs).logits
        probs = torch.softmax(logits, dim=1)  # [is_next, not_next]
        return probs[0, 0].item() >= threshold

# NSP 切片逻辑
def split_text_by_nsp(text, threshold=0.5):
    sentences = sent_tokenize(text)
    if len(sentences) <= 1:
        return [text]

    chunks = []
    buffer = [sentences[0]]

    for i in range(1, len(sentences)):
        if is_next_sentence(buffer[-1], sentences[i], threshold):
            buffer.append(sentences[i])
        else:
            chunks.append(" ".join(buffer))
            buffer = [sentences[i]]
    
    if buffer:
        chunks.append(" ".join(buffer))

    return chunks

# 示例文本
text = """
Artificial intelligence is transforming industries. Machines can now perform tasks like diagnosing diseases. 
However, many experts worry about job loss. Reskilling workers is becoming essential. 
AI was first proposed in the 1950s. The early expectations were overly optimistic.
"""

# 切片结果
chunks = split_text_by_nsp(text, threshold=0.6)
for i, chunk in enumerate(chunks):
    print(f"\nChunk {i+1}:\n{chunk}")

NSP 切片的优势

1. 语义感知的边界

NSP 模型能判断句子间是否语义连贯,避免误切语义相关的句群。

2. 切片内容更自然

每个 chunk 更像"自然段",对语言模型更友好,也更容易被 embedding 捕捉到。

3. 易于集成到 RAG pipeline

只需在原始分句之后添加 NSP 判定逻辑,即可作为 vectorization 之前的预处理。

潜在问题及优化建议

信息密度不均(Inconsistent information density)

有的 chunk 很短但内容贫乏,有的 chunk 很长且信息密集。解决方法:

  • 控制最小/最大句数限制

  • 对 chunk 加入关键词打分过滤

NSP 偶尔判断失误

NSP 本质是二分类,有一定误判概率。可以通过:

  • 提高阈值(更严格地判定"非连贯")

  • 多模型投票判定

实际应用场景

  • 构建 RAG 系统(文档问答、法律问答、产品说明 QA)

  • AI 助手的文档摘要模块

  • 信息提取、段落重组任务

总结

使用 BERT NSP 实现文本切片是一种兼顾语义完整性和实现简便性的优秀方法,特别适合构建高质量的文档检索系统(如 RAG)。它避免了固定窗口切片的语义割裂问题,生成的 chunk 更自然、上下文更丰富。

搭配向量搜索、chunk scoring、或 reranking 模块,将进一步提高整体系统的效果。