引言
在自然语言处理(NLP)任务中,数据质量是模型性能的关键因素之一。重复或冗余的数据会导致模型过度拟合或浪费计算资源,特别是在大语言模型(如 BERT、GPT 系列等)训练和推理阶段。传统的基于字符匹配的去重方法(如字符串哈希或编辑距离)在面对语义相似的文本时表现有限,而语义相似度算法则能更好地捕获文本之间的深层语义关系。
本文将介绍一种基于语义表示的去重方法,通过大语言模型生成的嵌入向量结合高效的相似度计算工具(如 FAISS),对大规模文本数据进行去重。此方法不仅适用于数据清洗,还可以应用在搜索引擎、推荐系统等需要衡量语义相似度的场景。
原理与方法
1. 传统去重方法的局限性
在 NLP 任务中,传统的去重方法包括:
- 字符串哈希:
基于文本的哈希值进行判重,适合完全重复的文本,但无法处理语义相似但表达不同的情况,例如:- 文本 A:我喜欢吃苹果。
- 文本 B:苹果是我最喜欢的水果。
虽然两者语义相近,但哈希值完全不同。
- 编辑距离(Levenshtein Distance):
衡量两个字符串的编辑代价,适合处理少量字符差异的文本,但无法捕捉深层语义关系。
上述方法对文本的语义相似性缺乏鲁棒性,特别是在短文本或同义表达常见的场景下。例如,问答生成、文档去重、语料清洗等任务中,语义相似的重复数据可能会严重影响模型性能。
2. 基于语义嵌入的去重
语义嵌入(Semantic Embedding)是一种将文本映射到高维向量空间的技术,向量的物理距离或角度可以反映文本语义的相似程度。常见的嵌入生成模型包括:
- BERT、RoBERTa、GPT 等大语言模型:能够生成上下文相关的语义表示。
- Sentence-BERT(SBERT):专为语义相似度任务设计,提升了嵌入的语义表达能力。
基本流程:
-
- 文本嵌入生成:
使用大语言模型将文本转化为固定维度的向量表示(如 768 维)。
- 文本嵌入生成:
-
- 相似度计算:
通过数学距离(如余弦相似度或内积)衡量文本向量之间的相似性。
- 相似度计算:
-
- 去重判断:
基于相似度阈值判断文本是否为重复内容。
- 去重判断:
3. 相似度计算方法对比
在语义嵌入的基础上,常用的相似度计算方法包括:
3.1. 余弦相似度(Cosine Similarity)
余弦相似度衡量两个向量的夹角余弦值,范围为 [ − 1 , 1 ] [-1, 1] [−1,1],归一化后范围为 [ 0 , 1 ] [0, 1] [0,1]。公式如下:
Cosine Similarity ( A , B ) = A ⋅ B ∥ A ∥ ∥ B ∥ \text{Cosine Similarity}(A, B) = \frac{A \cdot B}{\|A\| \|B\|} Cosine Similarity(A,B)=∥A∥∥B∥A⋅B
- 优点:消除向量模长的影响,只关注向量方向。
- 缺点:计算开销稍高。
3.2. 内积相似度(Inner Product Similarity)
内积相似度直接计算两向量的点积值:
Inner Product ( A , B ) = A ⋅ B \text{Inner Product}(A, B) = A \cdot B Inner Product(A,B)=A⋅B
- 优点:计算简单,速度快。
- 缺点:受向量模长影响,需要确保输入向量已归一化(模长为 1),否则结果不等价于余弦相似度。
欧几里得距离(Euclidean Distance)
衡量两个向量在高维空间中的直线距离:
Euclidean Distance ( A , B ) = ∑ i = 1 n ( A i − B i ) 2 \text{Euclidean Distance}(A, B) = \sqrt{\sum_{i=1}^n (A_i - B_i)^2} Euclidean Distance(A,B)=i=1∑n(Ai−Bi)2
- 优点:适合绝对位置相关的任务。
- 缺点:不适合捕获方向性的语义相似度。
4. 高效的大规模相似度计算
直接比较所有嵌入向量的相似度在大规模数据中效率低下(复杂度为 O ( n 2 ) O(n^2) O(n2))。为此,我们借助 FAISS(Facebook AI Similarity Search)工具,能够在百万级甚至亿级数据中高效实现近似最近邻搜索。
4.1. FAISS 简介
FAISS 是一个高效的相似度搜索库,专为高维向量的最近邻搜索设计,支持以下特性:
- 多种索引结构:
- Flat:暴力搜索,适合中小规模数据。
- IVF(倒排文件索引):适合大规模数据。
- PQ(分组量化):进一步压缩内存占用。
- GPU 加速:支持 GPU 版本,在大规模数据上极大提升搜索速度。
- 灵活的距离度量:支持内积、余弦、欧几里得距离等。
4.2. 使用 FAISS 的语义去重流程
- 初始化 FAISS 索引:选择适合任务的数据结构(如 IndexFlatIP)。
- 添加向量:将嵌入向量添加到索引。
- 查询相似度:对每个新向量,查找与索引中最近的向量,判断是否重复。
代码实现
import json
from transformers import BertTokenizer, BertModel
import torch
from tqdm import tqdm
import faiss
from typing import List, Dict, Union
class TextDeduplicatorWithFAISS:
"""
使用 FAISS 索引实现的文本去重类(基于余弦相似度)。
"""
def __init__(self, model_name: str = 'bert-base-chinese', device: str = None) -> None:
"""
初始化文本去重类。
参数:
- model_name: 使用的预训练模型名称,默认为 'bert-base-chinese'。
- device: 指定运行设备('cpu' 或 'cuda'),默认为自动检测。
"""
self.tokenizer = BertTokenizer.from_pretrained(model_name)
self.model = BertModel.from_pretrained(model_name)
self.device = device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
self.model = self.model.to(self.device)
# 初始化 FAISS 索引
self.embedding_dim = 768 # BERT 输出嵌入维度
self.index = faiss.IndexFlatIP(self.embedding_dim) # 使用内积(IP)作为相似度度量
self.index_ids = [] # 存储对应嵌入的 ID,方便后续处理
def get_embeddings(self, texts: List[str]) -> torch.Tensor:
"""
计算文本的嵌入表示,并进行归一化。
参数:
- texts: 要计算嵌入的一组文本列表。
返回:
- 归一化后的文本嵌入张量,形状为 (batch_size, hidden_size)。
"""
inputs = self.tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=512)
inputs = inputs.to(self.device) # 将输入张量移动到指定设备
with torch.no_grad(): # 禁用梯度计算以节省内存
outputs = self.model(**inputs)
embeddings = outputs.last_hidden_state[:, 0, :].cpu() # 获取 [CLS] 的嵌入并移动到 CPU
# 对嵌入进行归一化处理(实现余弦相似度)
embeddings = embeddings / torch.norm(embeddings, dim=1, keepdim=True)
return embeddings
def is_duplicate(self, embedding: torch.Tensor, threshold: float = 0.9) -> bool:
"""
检查一个嵌入是否与 FAISS 索引中的嵌入重复。
参数:
- embedding: 待检查的嵌入向量,形状为 (1, hidden_size)。
- threshold: 相似度的阈值,默认为 0.9。
返回:
- 是否为重复项(True / False)。
"""
if self.index.ntotal == 0: # 如果索引为空,肯定不是重复
return False
# 通过 FAISS 查找最近的向量及其相似度
embedding_np = embedding.numpy() # 转为 NumPy 格式
distances, _ = self.index.search(embedding_np, k=1) # 查找最近的 1 个向量
# 检查最近向量的相似度是否高于阈值
max_similarity = distances[0][0] # FAISS 返回的是归一化向量的内积(等价于余弦相似度)
return max_similarity >= threshold
def add_to_index(self, embedding: torch.Tensor, doc_id: int) -> None:
"""
将新的嵌入添加到 FAISS 索引中。
参数:
- embedding: 要添加的嵌入向量,形状为 (1, hidden_size)。
- doc_id: 该嵌入对应的文档 ID。
"""
embedding_np = embedding.numpy() # 转为 NumPy 格式
self.index.add(embedding_np) # 添加到索引中
self.index_ids.append(doc_id) # 保存对应的文档 ID
def process_and_save(self, input_path: str, output_path: str, threshold: float = 0.9) -> None:
"""
处理输入文件,去除相似文本并保存到输出文件。
参数:
- input_path: 输入 JSONL 文件路径。
- output_path: 输出 JSONL 文件路径。
- threshold: 去重的相似度阈值,默认值为 0.9。
"""
doc_id = 0 # 用于标记每条文档的唯一 ID
with open(input_path, 'r', encoding='utf-8') as infile, open(output_path, 'w', encoding='utf-8') as outfile:
for line in tqdm(infile, desc="Processing lines"):
item: Dict[str, Union[str, int, float]] = json.loads(line) # 从 JSONL 文件中读取一条数据
output_text: str = item['output'] # 获取文本内容
# 获取当前文本的嵌入
current_embedding = self.get_embeddings([output_text])
# 检查是否为重复
if not self.is_duplicate(current_embedding, threshold):
# 如果不重复,保存文本,并将嵌入添加到索引
outfile.write(json.dumps(item, ensure_ascii=False) + '\n')
self.add_to_index(current_embedding, doc_id)
doc_id += 1
# 使用示例
if __name__ == "__main__":
# 初始化去重器
deduplicator = TextDeduplicatorWithFAISS(model_name='bert-base-chinese')
# 去重并保存结果
deduplicator.process_and_save(
input_path='=./processed_unique_data-5.jsonl',
output_path='=./processed_unique_data-6.jsonl',
threshold=0.95
)
数据示例:
{"id": 1, "output": "什么是人工智能?人工智能是指让机器具备人类智能的技术。"}
{"id": 2, "output": "人工智能的定义是什么?人工智能是赋予机器类似人类智能的能力。"}
总结
本文介绍了一种基于语义嵌入的大规模文本去重方法,通过结合大语言模型(如 BERT)和高效相似度搜索工具(FAISS),实现了对语料库的语义级去重。该方法具有以下优点:
- 高精度:捕捉语义相似性,避免遗漏同义表达的重复数据。
- 高扩展性:支持大规模数据处理,适用于百万级文本的去重任务。
- 通用性强:不仅适用于去重,还可扩展至相似文本检索、推荐系统等任务。