我有一个由大约 500 个不同段落组成的数据集。对于每个段落，我都在尝试查看是否有任何其他段落的链接。基于此，我创建了段落对。我以前尝试将此问题作为二元问题（0 或 1，是否存在链接）来处理，但我现在想尝试排名（为每个段落对分配概率）。
我的问题是：如何随机拆分测试集和训练集，但将每个段落的所有段落对保留在同一集中？例如，对于第 1 段，我想要测试集中或训练集中的所有关联对（1-2、1-3、1-4、1-5...1-500）。例如，如果一半对在训练集中，我的排名将不起作用，因为那时测试集的排名将丢失一些对......
格式
A段| B段| 标签 | 特征...
第 1 段 | 第 4 段 | 1 | ...
第 2 段 | 第 6 段 | 1 | ...
第 6 段 | 第 8 段 | 0 | ...
第 10 段 | 第 2 段 | 1 | ...
我正在使用 sklearn train_test_split：
import pandas as pd
from sklearn.model_selection import train_test_split
feature_headers = ['tfidf_cosine', 'count_vec_cosine', 'lda_50topics_cosine', 'lda_200topics_cosine']
target_header = ['label']
train_x, test_x, train_y, test_y = train_test_split(result[feature_headers], result[target_header],
train_size=0.7)