当前位置：首页 > news >正文

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

news 2025/8/3 16:53:23

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

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

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

什么是 NSP？

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

“B 是不是 A 的下一句？”

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

如何用 NSP 做文本切片？

步骤概览：

将文本分句（sentence tokenization）
对每对相邻句子使用 BERT NSP 判断是否连贯
如果不连贯 → 作为一个切片边界
合并成语义一致的 chunk，用于向量化和检索

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).logitsprobs = 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}")