如何根据组 ID 生成训练测试拆分？

Question

我有以下数据：

pd.DataFrame({'Group_ID':[1,1,1,2,2,2,3,4,5,5],
          'Item_id':[1,2,3,4,5,6,7,8,9,10],
          'Target': [0,0,1,0,1,1,0,0,0,1]})

   Group_ID Item_id Target
0         1       1      0
1         1       2      0
2         1       3      1
3         2       4      0
4         2       5      1
5         2       6      1
6         3       7      0
7         4       8      0
8         5       9      0
9         5      10      1

我需要根据“Group_ID”将数据集拆分为训练集和测试集，以便 80% 的数据进入训练集，20% 进入测试集。

也就是说，我需要我的训练集看起来像：

    Group_ID Item_id Target
0          1       1      0
1          1       2      0
2          1       3      1
3          2       4      0
4          2       5      1
5          2       6      1
6          3       7      0
7          4       8      0

和测试集：

Test Set
   Group_ID Item_id Target
8         5       9      0
9         5      10      1

什么是最简单的方法来做到这一点？据我所知，sklearn 中的标准 test_train_split 函数不支持按组拆分，我也可以指示拆分的大小（例如 80/20）。

Answer 1

我想出了答案。这似乎有效：

train_inds, test_inds = next(GroupShuffleSplit(test_size=.20, n_splits=2, random_state = 7).split(df, groups=df['Group_Id']))

train = df.iloc[train_inds]
test = df.iloc[test_inds]