cluster_horserace.py

from sklearn.cluster import AffinityPropagation, KMeans, MiniBatchKMeans
from sklearn.feature_extraction.text import TfidfVectorizer
import odcdata
import pprint as pp

q = "select UPPER(description) as description  from odc_bank_transaction limit 1000"
data = odcdata.read_gp_as_dict(q)
descriptions = [row['description'] for row in data]

tfidf = TfidfVectorizer()
X = tfidf.fit_transform(descriptions)


cluterers = {
    "affinity": AffinityPropagation(),
    "kmeans": KMeans(n_clusters=50),
    "mini-batch-kmeans": MiniBatchKMeans(n_clusters=50)
}

for name, clusterer in cluterers.iteritems():
    print "*"*80
    print "fitting %s" % name
    clusterer.fit(X)
    clusters = {}
    for desc, label in zip(descriptions, clusterer.labels_):
        clusters[label] = clusters.get(label, [])
        clusters[label] += [desc]

    raw_input("ready to print results for %s?" % name)
    pp.pprint(clusters)
    raw_input("ready for next clusterer?")