如何加速Pandas处理5000万行Spacy句子相似度计算任务？

尝试用Spacy计算两组问题的句子相似度，再通过Pandas输出包含3列的DataFrame（样本问题、用户输入、相似度分数），预计生成约5000万行数据，已运行12小时仍未完成。

小数据集测试时代码正常，但大数据集处理极慢，相关代码如下：

import spacy
import pandas as pd

# 假设已加载nlp模型，例如 nlp = spacy.load("de_core_news_md")
user_inputs_vector = [nlp(row) for row in user_inputs_df["User_Input"]] # 向量化用户输入
sample_df_vector = [nlp(row) for row in sample_df["Frage"]] # 向量化样本句子

# 嵌套循环计算相似度并收集结果
similarity_score_list = [] 
sample_list = [] 
user_input_list = []

for i in range (len(sample_df_vector)):
    for j in range (len(user_inputs_vector)):    
        similar_frage = user_inputs_vector[j].similarity(sample_df_vector[i])
        similarity_score_list.append(similar_frage)
        sample_list.append(sample_df_vector[i])
        user_input_list.append(user_inputs_vector[j])
        similarity_dataframe = pd.DataFrame(list(zip(sample_list, user_input_list, similarity_score_list)), columns = ["Samples", "User Inputs", "Similarity Score"])

1. 移除循环内的DataFrame创建 #

原代码每次循环都重新生成DataFrame，反复触发内存分配与数据拷贝，是核心性能浪费点。将DataFrame创建移到循环结束后：

similarity_score_list = [] 
sample_list = [] 
user_input_list = []

for i in range(len(sample_df_vector)):
    sample_doc = sample_df_vector[i]
    sample_text = sample_doc.text  # 存储文本而非Doc对象，大幅节省内存
    for j in range(len(user_inputs_vector)):    
        user_doc = user_inputs_vector[j]
        similar_frage = user_doc.similarity(sample_doc)
        similarity_score_list.append(similar_frage)
        sample_list.append(sample_text)
        user_input_list.append(user_doc.text)

# 循环结束后一次性创建DataFrame
similarity_dataframe = pd.DataFrame(
    list(zip(sample_list, user_input_list, similarity_score_list)),
    columns=["Samples", "User Inputs", "Similarity Score"]
)

2. 用矩阵运算替代嵌套循环（核心优化） #

Spacy的Doc对象可通过.vector提取向量，转为numpy矩阵后用矩阵乘法一次性计算所有相似度，时间复杂度从O(NM)降至O(NK + K*M)（K为向量维度），速度提升几个数量级：

import numpy as np

# 提取所有样本和用户输入的向量，转为numpy矩阵
sample_vectors = np.array([doc.vector for doc in sample_df_vector])
user_vectors = np.array([doc.vector for doc in user_inputs_vector])

# 计算余弦相似度：(用户向量 @ 样本向量转置) / (用户向量模长 @ 样本向量模长转置)
user_norms = np.linalg.norm(user_vectors, axis=1, keepdims=True)
sample_norms = np.linalg.norm(sample_vectors, axis=1, keepdims=True)
similarity_matrix = (user_vectors @ sample_vectors.T) / (user_norms @ sample_norms.T)

# 将矩阵转为长格式DataFrame
sample_texts = sample_df["Frage"].tolist()
user_texts = user_inputs_df["User_Input"].tolist()

# 生成所有配对的索引并扁平化
sample_indices, user_indices = np.meshgrid(range(len(sample_texts)), range(len(user_texts)))

similarity_dataframe = pd.DataFrame({
    "Samples": np.array(sample_texts)[sample_indices.flatten()],
    "User Inputs": np.array(user_texts)[user_indices.flatten()],
    "Similarity Score": similarity_matrix.flatten()
})

3. 内存优化：分块处理或直接写CSV #

5000万行数据会占用大量内存（按每行100字节计算约5GB），若内存不足，建议分块处理并直接写入CSV，避免一次性加载所有数据：

chunk_size = 1000  # 根据内存容量调整块大小
with open("similarity_results.csv", "w", encoding="utf-8") as f:
    # 写入表头
    f.write("Samples,User Inputs,Similarity Score\n")
    for i in range(0, len(sample_vectors), chunk_size):
        sample_chunk = sample_vectors[i:i+chunk_size]
        sample_text_chunk = sample_texts[i:i+chunk_size]
        # 计算当前块的相似度矩阵
        chunk_norms = np.linalg.norm(sample_chunk, axis=1, keepdims=True)
        chunk_similarity = (user_vectors @ sample_chunk.T) / (user_norms @ chunk_norms.T)
        # 生成当前块的配对数据并写入文件
        chunk_sample_indices, chunk_user_indices = np.meshgrid(range(len(sample_text_chunk)), range(len(user_texts)))
        for s_idx, u_idx, score in zip(chunk_sample_indices.flatten(), chunk_user_indices.flatten(), chunk_similarity.flatten()):
            # 转义引号避免CSV格式错误
            s_text = sample_text_chunk[s_idx].replace('"', '""')
            u_text = user_texts[u_idx].replace('"', '""')
            f.write(f'"{s_text}","{u_text}",{score}\n')

4. 优化文本向量化过程 #

使用Spacy的nlp.pipe批量处理文本，比逐个调用nlp(row)更快，还可开启多进程加速：

# batch_size：每次批量处理的文本数量；n_process=-1：使用所有CPU核心
user_inputs_vector = list(nlp.pipe(user_inputs_df["User_Input"], batch_size=512, n_process=-1))
sample_df_vector = list(nlp.pipe(sample_df["Frage"], batch_size=512, n_process=-1))

内容的提问来源于stack exchange，提问作者klam