Datasetpublic

SetFit/enron_spam

This is a version of the Enron Spam Email Dataset, containing emails (subject + message) and a label whether it is spam or ham.

sourceHugging Faceupdated 5y agoView on Hugging Face

21likes6.4kdownloads

rootprepare.py

prepare.py38 linesDownload Raw Back to root

1import pandas as pd
2from collections import Counter
3import json
4import random
5
6df = pd.read_csv("enron_spam_data.csv")
7df.fillna('', inplace=True)
8print(df)
9label2id = {'ham': 0, 'spam': 1}
10
11rows = [{'message_id': row['Message ID'], 
12         'text': (row['Subject']+" "+row['Message']).strip(),
13         'label': label2id[row['Spam/Ham']],
14         'label_text': row['Spam/Ham'],
15         'subject': row['Subject'],
16         'message': row['Message'],
17         'date': row['Date']
18         } for idx, row in df.iterrows()]
19
20random.seed(42)
21random.shuffle(rows)
22
23num_test = 2000
24splits = {'test': rows[0:num_test], 'train': rows[num_test:]}
25
26print("Train:", len(splits['train']))
27print("Test:", len(splits['test']))
28
29num_spam = Counter()
30
31for row in splits['test']:
32    num_spam[row['label']] += 1
33print(num_spam)
34
35for split in ['train', 'test']:
36    with open(f'{split}.jsonl', 'w') as fOut:
37        for row in splits[split]:
38            fOut.write(json.dumps(row)+"\n")