Deduplicating Data with Clustering and Similarity

I’ve been preparing a conversation dataset lately and found quite a few similar examples in the base data. Time to get rid of those duplicates.

First, I converted every example into a vector using the BGE-M3 embedding model.

You can read about the selection here. The main reason was that BGE-M3 was the only embedding model I could find at the time that supported multiple languages, including Chinese.

Next, I used DBSCAN to gather similar examples into the same cluster.

For each cluster, I calculated the pairwise cosine similarity between all its examples. If an example’s highest similarity to another example exceeded the threshold, I removed it. That way, the remaining examples would have relatively low similarity to one another.

Finally, I saved the cleaned data to a new file for later use.

I also made a few optimizations, such as asynchronous requests and multiprocessing, to speed things up. Progress bars and logging made the output easier to follow, too.

I’ll release this as part of a dataset toolchain!

GitHub: https://github.com/Ce-daros/Collider

Code snippets

1. Read the input data and concatenate the system message and conversation:

读取数据并拼接 system 和 conversation  
logging.info("Reading and concatenating data...")
all_texts = []
with open(input_file, 'r', encoding='utf-8') as f:
    data = json.load(f)
    logging.info(f"Loaded {len(data)} entries from {input_file}")
    for entry in data:
        system = entry.get("system", "")
        conversation = entry.get("conversations", [])
        if conversation:
            value = conversation[0].get("value", "")
            text = system + value
            all_texts.append(text)

2. Calculate embeddings:

logging.info("Calculating embeddings...")

all_embeddings = []
batch_count = 0
with torch.no_grad():
    for i in tqdm(range(0, len(all_texts), batch_size), desc="Batches"):
        batch = all_texts[i:i+batch_size]
        # 使用模型计算嵌入向量
        inputs = tokenizer(batch, padding=True, truncation=True, return_tensors="pt").to(device)
        output = model(**inputs, return_dict=True)
        dense_output = output.dense_output
        embeddings = dense_output.cpu().numpy()
        all_embeddings.append(embeddings)

3. DBSCAN clustering and similarity filtering:

执行 DBSCAN 聚类
logging.info("Performing DBSCAN clustering...")
clustering = DBSCAN(eps=eps, min_samples=min_samples, metric='cosine')
cluster_labels = clustering.fit_predict(list(tqdm(all_embeddings, desc="DBSCAN input")))  # 将 tqdm 迭代器转换为列表

# 移除高相似度对
if True:
    logging.info("Removing similar vectors using DBSCAN clustering and similarity threshold...")
    unique_data = []
    discarded_vectors = set()  # 存储已舍弃的向量(字符串形式)
    for label in np.unique(cluster_labels):
        if label != -1:
            # 对于每个簇
            cluster_indices = np.where(cluster_labels == label)[0]
            cluster_embeddings = all_embeddings[cluster_indices]
            cluster_data = [data[i] for i in cluster_indices]
            
            # 计算簇内相似度矩阵
            similarity_matrix = cosine_similarity(cluster_embeddings)
            
            # 根据相似度阈值选择代表
            representatives = []
            for i in range(len(cluster_data)):
                similar_indices = np.where(similarity_matrix[i] > similarity_threshold)[0]
                if len(similar_indices) == 1:
                    # 只有一个相似向量, 保留该向量
                    vector_str = str(cluster_data[i])
                    if vector_str not in discarded_vectors:
                        representatives.append(cluster_data[i])
                    else:
                        discarded_vectors.add(vector_str)
                else:
                    # 有多个相似向量, 选择第一个作为代表
                    representative_index = similar_indices[0]
                    if i == representative_index:
                        vector_str = str(cluster_data[i])
                        if vector_str not in discarded_vectors:
                            representatives.append(cluster_data[i])
                    else:
                        discarded_vectors.add(str(cluster_data[i]))
                        
            unique_data.extend(representatives)
        else:
            # 对于噪声点直接保留
            noise_indices = np.where(cluster_labels == -1)[0]
            unique_data.extend([data[i] for i in noise_indices if str(data[i]) not in discarded_vectors])
            
    logging.info(f"Saving {len(unique_data)} entries to {output_file}")
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(unique_data, f, ensure_ascii=False, indent=4)
        
    logging.info(f"After removing: {len(unique_data)}. Before removing: {len(data)}.")

An entire afternoon of wrestling with AI, plus a whole night… and the first module is finally done. Here’s the rough dataset-generation workflow: