# clustering.py import json import logging import numpy as np from sklearn.cluster import KMeans from sklearn.metrics import silhouette_score from tracescope.analysis.metric_adapter import learn_metric, safe_transform # ── Global logger setup ── logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s:%(message)s", datefmt="%Y-%m-%d %H:%M:%S") logger = logging.getLogger(__name__) def cluster_embeddings(embedding_list, n_clusters=None, user_clusters=None, min_k=3, random_state=42): # build & normalize matrix = np.array(embedding_list, dtype=float) matrix = matrix / np.linalg.norm(matrix, axis=1, keepdims=True) n_samples = matrix.shape[0] logger.info(f"cluster_embeddings ▶ n_samples={n_samples}, requested n_clusters={n_clusters}, random_state={random_state}") if n_samples < 2: raise ValueError("Need at least 2 samples for clustering") # auto‑select k by silhouette if n_clusters is None: max_k = min(10, n_samples - 1) if user_clusters: # 2-stage grid: α ∈ {0,.25,.5,.75,1} and k ∈ [2…max_k] matrix_orig = matrix.copy() L = learn_metric(matrix_orig, user_clusters) best_score = -1 best_k = 2 best_alpha = 0 best_labels = None best_centroids = None for alpha in (0, 0.25, 0.5, 0.75, 1): Xα = safe_transform(matrix_orig, L, alpha) for k in range(2, max_k + 1): km = KMeans(n_clusters=k, init='k-means++', random_state=random_state) lbl = km.fit_predict(Xα) sil = silhouette_score(Xα, lbl) if sil > best_score: best_score, best_k = sil, k best_alpha = alpha best_labels = lbl best_centroids = km.cluster_centers_ logger.info( f"user feedback → α={best_alpha}, k={best_k}, silhouette={best_score:.4f}" ) n_clusters = best_k labels = best_labels centroids = best_centroids else: # no feedback → original silhouette-only loop best_score = -1 best_k = min_k best_labels = None best_centroids = None for k in range(min_k, max_k + 1): logger.info(f"Trying k={k}") km = KMeans(n_clusters=k, init='k-means++', random_state=random_state) lbl = km.fit_predict(matrix) sil = silhouette_score(matrix, lbl) logger.info(f" k={k} silhouette={sil:.4f}") if sil > best_score: best_score, best_k = sil, k best_labels = lbl best_centroids = km.cluster_centers_ # unbalanced 2-cluster guard as before counts = np.bincount(best_labels) if best_k == 2 and (counts.min() < 2 or counts.max()/counts.min() > 10): logger.info("Detected unbalanced 2-cluster; searching k=3…max_k") alt_best = -1 alt_labels = None alt_k = None for k in range(3, min(max_k + 1, 6)): lbl = KMeans(n_clusters=k, random_state=random_state).fit_predict(matrix) s = silhouette_score(matrix, lbl) logger.info(f" alt_k={k} silhouette={s:.4f}") if s > alt_best: alt_best, alt_k, alt_labels = s, k, lbl if alt_k: logger.info(f"Switching to k={alt_k} with silhouette={alt_best:.4f}") best_k, best_labels = alt_k, alt_labels best_centroids = KMeans(n_clusters=alt_k, random_state=random_state).fit(matrix).cluster_centers_ n_clusters = best_k labels = best_labels centroids = best_centroids else: kmeans = KMeans(n_clusters=n_clusters, init='k-means++', random_state=random_state) labels = kmeans.fit_predict(matrix) centroids = kmeans.cluster_centers_ # build output clusters = [] for i in range(n_clusters): idxs = [int(idx) for idx, lbl in enumerate(labels) if lbl == i] clusters.append({ "cluster_id": i, "indices": idxs, "centroid": centroids[i].tolist() }) return { "n_clusters": n_clusters, "clusters": clusters, "labels": labels.tolist(), } def main(embeddings_json, n_clusters=None, aspect=None, user_clusters=None, compute_moving=False, min_k=3, random_state=42): # Load all embeddings directly from JSON embedding_list = json.loads(embeddings_json) # 2) if user_clusters came in as a JSON string, decode it, # and make sure every index is an int # ──────────────────────────────────────────────────────────────────────── # sanitize user_clusters: if it’s a JSON string, parse it; then # coerce absolutely every index to a Python int, or error. if user_clusters: # 1) parse JSON if we were handed a string if isinstance(user_clusters, str): try: user_clusters = json.loads(user_clusters) except Exception as e: raise ValueError(f"Could not parse user_clusters JSON: {e}") # 2) coerce every element to int try: user_clusters = [ [int(idx) for idx in cluster] for cluster in user_clusters ] except Exception as e: raise ValueError( f"Invalid user_clusters format (must be List[List[int]]): {user_clusters}" ) from e # ──────────────────────────────────────────────────────────────────────── # Cluster them (no deduplication or pruning) result = cluster_embeddings( embedding_list, n_clusters, user_clusters=user_clusters, min_k=min_k, random_state=random_state, ) # print("RESULT:", result) # ── Optional moving‐average “closeness” to each cluster ── if compute_moving: try: # rebuild & normalize matrix matrix = np.array(embedding_list, dtype=float) matrix = matrix / np.linalg.norm(matrix, axis=1, keepdims=True) n = matrix.shape[0] # weights decay with distance in list: w_{i,j} = 1 / (1 + |i-j|) idx = np.arange(n) W = 1 / (1 + np.abs(idx[:, None] - idx[None, :])) # compute weighted‐average embedding for each position weighted = (W[:, :, None] * matrix[None, :, :]).sum(axis=1) / W.sum(axis=1)[:, None] # extract centroids from the clustering result centroids = np.array([c["centroid"] for c in result["clusters"]]) # for each position, compute similarity = 1/(1 + distance) to each centroid moving_sim = ( 1 / ( 1 + np.linalg.norm( centroids[None, :, :] - weighted[:, None, :], axis=2 ) ) ).tolist() result["moving_similarity"] = moving_sim except Exception as e: logger.error(f"Error computing moving‐average similarity: {e}") # Return the raw clustering result as JSON return json.dumps(result)