Как сгенерировать поезд-тест-сплит на основе идентификатора группы?

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% в тестовом наборе.

То есть мне нужно, чтобы мой тренировочный набор выглядел примерно так:

Training Set:       
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

Какой самый простой способ сделать это?Насколько я знаю, стандартная функция test_train_split в sklearn не поддерживает разбиение по группам, так что я также могу указать размер разделения (например, 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]