fix KMeans arguments to support older sklearn versions

Helsinki-NLP · Sep 6, 2023 · e8a3f4e · e8a3f4e
1 parent 0eef64f
commit e8a3f4e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/opusfilter/autogen_cluster.py b/opusfilter/autogen_cluster.py
@@ -37,7 +37,7 @@ def __init__(self, score_file, n=2):
         self.standard_data = self.scaler.fit_transform(self.df)
 
         logger.info('Training KMeans with %s clusters', n)
-        self.kmeans = KMeans(n_clusters=n, random_state=0, n_init='auto').fit(self.standard_data)
+        self.kmeans = KMeans(n_clusters=n, random_state=0, init='k-means++', n_init=1).fit(self.standard_data)
         self.labels = self.kmeans.labels_
         self.cluster_centers = self.scaler.inverse_transform(self.kmeans.cluster_centers_)
         self._noisy_label = self._get_noisy_label()