Upload clustering_utils.py

clustering_utils.py

@@ -1,16 +1,21 @@
 from sentence_transformers import SentenceTransformer
 import hdbscan
 from sklearn.metrics import silhouette_score, davies_bouldin_score
+import numpy as np
 
 model = SentenceTransformer("shibing624/text2vec-bge-large-chinese")
 
 def cluster_sentences(sentences):
-    embeddings = model.encode(sentences, normalize_embeddings=True)
-    clusterer = hdbscan.HDBSCAN(min_cluster_size=3, prediction_data=True)
+    embeddings = model.encode(sentences)
+    clusterer = hdbscan.HDBSCAN(min_cluster_size=2, metric='euclidean')
     labels = clusterer.fit_predict(embeddings)
-    try:
-        sil = silhouette_score(embeddings, labels) if len(set(labels)) > 1 else -1
-        db = davies_bouldin_score(embeddings, labels) if len(set(labels)) > 1 else -1
-    except Exception:
-        sil, db = -1, -1
-    return labels, embeddings, {"silhouette": sil, "db": db}
+
+    valid_idxs = labels != -1
+    if np.sum(valid_idxs) > 1:
+        silhouette = silhouette_score(embeddings[valid_idxs], labels[valid_idxs])
+        db = davies_bouldin_score(embeddings[valid_idxs], labels[valid_idxs])
+    else:
+        silhouette, db = -1, -1
+
+    return labels, embeddings, {"silhouette": silhouette, "db": db}
+