kmeans.py

from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score


sum_of_sq_distance = []
silhouette_score_list = []
K = range(2, 80)

## Loop through K and calculate silhouette score
for k in K:
  print("KMeans - Num Cluster - ", k)
  km = KMeans(n_clusters=k, random_state=0).fit(x_train_lsa)
  sum_of_sq_distance.append(km.inertia_)
  silhouette_score_list.append(silhouette_score(x_train_lsa, km.labels_))
  
  
# Plot sum of sq distance error
plt.plot(K, sum_of_sq_distance, 'bx-')
plt.xlabel('k')
plt.ylabel('Sum of Squared distances')
plt.title("Elbow Method")
plt.show()


# Plot sum of silhouette score
plt.plot(K, silhouette_score_list, 'bx-')