modified knn.py