KMeans python model and test data