File size: 1,493 Bytes
2d06dcc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | import logging
from utils.metrics import clustering_score
from sklearn.metrics import confusion_matrix
class KMManager:
def __init__(self, args, data, model, logger_name = 'Discovery'):
self.logger = logging.getLogger(logger_name)
self.emb_train, self.emb_test = model.set_model(args, data, 'glove')
self.num_labels = data.num_labels
self.test_y = data.dataloader.test_true_labels
def train(self, *args):
self.logger.info('K-Means does not need training...')
pass
def test(self, args, data, show=True):
self.logger.info('K-Means start...')
from sklearn.cluster import KMeans
km = KMeans(n_clusters=self.num_labels, n_jobs=-1, random_state = args.seed)
km.fit(self.emb_train)
self.logger.info('K-Means finished...')
y_pred = km.predict(self.emb_test)
y_true = self.test_y
test_results = clustering_score(y_true, y_pred)
cm = confusion_matrix(y_true, y_pred)
if show:
self.logger.info
self.logger.info("***** Test: Confusion Matrix *****")
self.logger.info("%s", str(cm))
self.logger.info("***** Test results *****")
for key in sorted(test_results.keys()):
self.logger.info(" %s = %s", key, str(test_results[key]))
test_results['y_true'] = y_true
test_results['y_pred'] = y_pred
return test_results
|