| """Theme discovery utilities based on review embeddings.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import numpy as np |
| from sklearn.cluster import KMeans |
|
|
|
|
| @dataclass(slots=True) |
| class ThemeDiscoveryResult: |
| """Cluster assignments and human-readable theme labels.""" |
|
|
| assignments: list[int] |
| top_terms_by_cluster: dict[int, list[str]] |
| representative_row_index_by_cluster: dict[int, int] |
|
|
|
|
| def discover_themes( |
| embeddings: np.ndarray, |
| tfidf_matrix, |
| feature_names: np.ndarray, |
| n_clusters: int, |
| top_terms_per_cluster: int, |
| seed: int, |
| ) -> ThemeDiscoveryResult: |
| """Cluster review embeddings and derive representative terms per cluster.""" |
| if len(embeddings) == 0: |
| return ThemeDiscoveryResult( |
| assignments=[], |
| top_terms_by_cluster={}, |
| representative_row_index_by_cluster={}, |
| ) |
|
|
| cluster_count = max(1, min(n_clusters, len(embeddings))) |
| model = KMeans(n_clusters=cluster_count, random_state=seed, n_init="auto") |
| assignments = model.fit_predict(embeddings) |
|
|
| top_terms_by_cluster: dict[int, list[str]] = {} |
| representative_row_index_by_cluster: dict[int, int] = {} |
| for cluster_id in range(cluster_count): |
| cluster_rows = np.where(assignments == cluster_id)[0] |
| cluster_matrix = tfidf_matrix[cluster_rows] |
| weights = np.asarray(cluster_matrix.mean(axis=0)).ravel() |
| ranked_indices = weights.argsort()[::-1] |
| terms = [ |
| str(feature_names[index]) |
| for index in ranked_indices[:top_terms_per_cluster] |
| if weights[index] > 0 |
| ] |
| top_terms_by_cluster[cluster_id] = terms or ["misc"] |
| cluster_vectors = embeddings[cluster_rows] |
| centroid = model.cluster_centers_[cluster_id] |
| distances = np.linalg.norm(cluster_vectors - centroid, axis=1) |
| representative_row_index_by_cluster[cluster_id] = int( |
| cluster_rows[int(np.argmin(distances))] |
| ) |
|
|
| return ThemeDiscoveryResult( |
| assignments=assignments.tolist(), |
| top_terms_by_cluster=top_terms_by_cluster, |
| representative_row_index_by_cluster=representative_row_index_by_cluster, |
| ) |
|
|
