按比例标签拆分pandas数据帧

Question

假设我有一个数据帧包含100个句子(50个垃圾邮件，50个不是垃圾邮件(
目标：我需要将它们拆分为训练：测试数据，比率为80:20
，这将是80个测试数据(40个垃圾邮件+40个非垃圾邮件(和20个测试数据(10个垃圾邮件+10个非垃圾(

注意：我用的是熊猫，我需要这些比例本身就是一个变量，这样我就可以把它改为

我在哪里：

import pandas as pd
df = pd.DataFrame({'sentence': {0: 'FU bro',
1: 'Well thats kinda cool',
2: 'Haha thats so funny',
3: 'cant u make somethin else mtfk',
4: 'what a shame'},
'label': {0: 'spam', 1: 'not spam', 2: 'not spam', 3: 'spam', 4: 'spam'}})
spam = df.loc[df['label']=='spam']
not_spam = df.loc[df['label']=='not spam']
print(spam)
print(not_spam)
#print(df.loc[df['label']=='not spam'].sum)

这是我的数据帧的标题，看起来像：

标签不是垃圾邮件不是垃圾邮件垃圾邮件

句子
FU bro	垃圾邮件
这有点酷
哈哈，太有趣了
你不能做其他东西吗
真是太可惜了

试试这个：

train = df.groupby('label').sample(frac=.8)
test = df.loc[df.index.difference(train.index)]

您可以使用DataFrame.sample()进行以下操作：

training_data_ratio = 0.8
train_spam = spam.sample(frac=training_data_ratio, random_state=0)
test_spam = spam.drop(train_spam.index)

同样，对于非垃圾邮件数据也是如此。

---

此外，如果您需要检查有多少条目是垃圾邮件而不是垃圾邮件，您可以使用value_counts:

>>> df.label.value_counts()
spam        3
not spam    2
Name: label, dtype: int64