如何划分训练/测试集并将同段落的所有配对保留在同一集合

你的核心需求很明确：必须保证每个段落的所有关联配对（不管它是作为Paragraph A还是B）都完全落在训练集或测试集里，绝对不能拆分——否则排序任务会因为缺失完整的配对信息而失效。直接用train_test_split按行随机划分肯定不行，因为它会把同一个段落的配对拆到两个集合里。下面给你两种可行的解决方案，从简单到严谨：

方案一：基于唯一段落ID的划分（简单直接） #

如果你的数据集里段落之间的关联比较松散，或者可以接受少量跨集合配对的存在，可以先对所有唯一段落ID进行划分，再筛选对应的配对数据：

步骤： #

1. 从数据集中提取所有唯一的段落ID（合并Paragraph A和Paragraph B列）
2. 对这些段落ID做训练/测试划分
3. 把所有包含训练集段落的配对归入训练集，仅包含测试集段落的配对归入测试集

代码实现： #

import pandas as pd
from sklearn.model_selection import train_test_split

# 假设你的数据集存储在result DataFrame中
# 提取所有唯一段落ID
all_paragraphs = pd.concat([result['Paragraph A'], result['Paragraph B']]).unique()

# 划分段落ID为训练集和测试集
train_paragraphs, test_paragraphs = train_test_split(all_paragraphs, train_size=0.7, random_state=42)

# 筛选训练集（所有包含训练段落的配对）和测试集（仅包含测试段落的配对）
train_data = result[(result['Paragraph A'].isin(train_paragraphs)) | (result['Paragraph B'].isin(train_paragraphs))]
test_data = result[(result['Paragraph A'].isin(test_paragraphs)) & (result['Paragraph B'].isin(test_paragraphs))]

# 拆分特征和标签
feature_headers = ['tfidf_cosine', 'count_vec_cosine', 'lda_50topics_cosine', 'lda_200topics_cosine']
target_header = ['label']

train_x = train_data[feature_headers]
train_y = train_data[target_header]
test_x = test_data[feature_headers]
test_y = test_data[target_header]

⚠️ 注意：这个方案可能会产生“跨集合配对”（比如一个训练段落和一个测试段落的配对），这类配对如果留在训练集，会导致测试段落的部分配对出现在训练集里，违反你“每个段落的所有配对在同一集合”的要求。如果你的数据集里段落关联紧密，更推荐下面的严谨方案。

方案二：基于连通分量的划分（严谨无冲突） #

如果段落之间通过配对形成了关联链（比如段落1和2配对，2和3配对，那么1、2、3属于同一个关联组），必须保证整个关联组的所有段落和配对都在同一个集合里，否则会出现配对拆分的冲突。这时候需要先识别所有关联组（连通分量），再对关联组进行划分：

步骤： #

1. 用段落和配对构建关联图（段落是节点，配对是边）
2. 找出图中的所有连通分量（相互关联的段落组）
3. 对连通分量做训练/测试划分
4. 把属于训练连通分量的所有配对归入训练集，测试连通分量的配对归入测试集

代码实现： #

import pandas as pd
from sklearn.model_selection import train_test_split
from itertools import chain
import networkx as nx

# 构建段落关联图
G = nx.Graph()
# 为每个配对添加一条边
for _, row in result.iterrows():
    G.add_edge(row['Paragraph A'], row['Paragraph B'])

# 找出所有连通分量（相互关联的段落组）
connected_components = list(nx.connected_components(G))

# 划分连通分量为训练集和测试集
train_components, test_components = train_test_split(connected_components, train_size=0.7, random_state=42)

# 将连通分量转换为段落ID集合
train_paragraphs = set(chain.from_iterable(train_components))
test_paragraphs = set(chain.from_iterable(test_components))

# 筛选训练集和测试集（确保配对的两个段落都在同一集合）
train_data = result[(result['Paragraph A'].isin(train_paragraphs)) & (result['Paragraph B'].isin(train_paragraphs))]
test_data = result[(result['Paragraph A'].isin(test_paragraphs)) & (result['Paragraph B'].isin(test_paragraphs))]

# 拆分特征和标签
feature_headers = ['tfidf_cosine', 'count_vec_cosine', 'lda_50topics_cosine', 'lda_200topics_cosine']
target_header = ['label']

train_x = train_data[feature_headers]
train_y = train_data[target_header]
test_x = test_data[feature_headers]
test_y = test_data[target_header]

说明：这个方案需要安装networkx库（执行pip install networkx即可），它能完美保证所有相互关联的段落都被分到同一个集合，完全避免配对拆分的问题，非常适合你的排序任务——模型训练时能看到某个段落的全部关联配对，测试时也能完整评估该段落的排序能力。

为什么你之前的方法不适用？ #

你之前直接对特征和标签用train_test_split，是按单个配对行随机拆分的。比如段落1的配对1-2可能被分到训练集，而1-3却被分到测试集——这就把同一个段落的配对拆得七零八落，排序任务需要的完整段落关联信息就缺失了，模型自然没法正常学习和评估。

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