| # File 8: scripts/cluster.py | |
| import torch | |
| import numpy as np | |
| from sklearn.cluster import KMeans | |
| from utils.data_loader import UserFeedbackDataset | |
| def update_latent_clusters(): | |
| dataset = UserFeedbackDataset() | |
| latents = torch.cat([data[0] for data in dataset]).numpy() | |
| # Cluster latent vectors | |
| kmeans = KMeans(n_clusters=10) | |
| clusters = kmeans.fit_predict(latents) | |
| # Update VAE prior with cluster centers | |
| cluster_centers = torch.tensor(kmeans.cluster_centers_, dtype=torch.float32) | |
| # Save new prior distribution | |
| torch.save(cluster_centers, 'storage/models/latent_prior.pt') |