Feature Extractor for Spam Filtering

spam_feature.py

import sys
import os
import gzip

def ngram(T, n):
    return ['%dgram=%s' % (n, '_'.join(T[i:i+n])) for i in range(len(T)-n+1)]

def process(fo, fi, label):
    F = []
    for line in fi:
        fields = line.strip('\n').split(' ')
        F += ngram(fields, 1)
        F += ngram(fields, 2)
    fo.write('%s ' % label)
    fo.write('%s' % ' '.join([f.replace(':', '__COLON__') for f in F]))
    fo.write('\n')

if __name__ == '__main__':
    fo = sys.stdout
    for src in sys.argv[1:]:
        label = '+1' if os.path.basename(src).startswith('spmsg') else '-1'
        fi = gzip.GzipFile(src) if src.endswith('.gz') else open(src)
        process(fo, fi, label)