Embedding Compression for Clustering

import os
import torch
import pickle
import random
import numpy as np
import torch.nn as nn

from tqdm import tqdm
from sklearn.datasets import fetch_20newsgroups
from sklearn.cluster import KMeans
from sklearn.metrics import normalized_mutual_info_score
from torch.utils.data import DataLoader, Dataset
from sklearn.model_selection import train_test_split
from sentence_transformers import SentenceTransformer

def set_seed(seed_value = 42):
  random.seed(seed_value)
  np.random.seed(seed_value)
  torch.manual_seed(seed_value)
  if torch.cuda.is_available():
      torch.cuda.manual_seed(seed_value)
      torch.cuda.manual_seed_all(seed_value)
      torch.backends.cudnn.deterministic = True
      torch.backends.cudnn.benchmark = False

set_seed()

class TextDataset(Dataset):
    def __init__(self, texts, labels, model_name='all-MiniLM-L6-v2'):
        self.texts = texts
        self.labels = labels
        self.model = SentenceTransformer(model_name)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        with torch.no_grad():
            emb = self.model.encode(text, convert_to_tensor=True)
        label = self.labels[idx]
        return emb, label

def load_and_process_split(split_name, model_name='all-MiniLM-L6-v2'):
    with open(f"{split_name}_texts.pkl", "rb") as f:
        texts, labels = pickle.load(f)

    dataset = TextDataset(texts, labels, model_name=model_name)
    loader = DataLoader(dataset, batch_size=32, shuffle=False)

    X = []
    y = []
    for emb, label in tqdm(loader):
        X.append(emb.cpu().numpy())
        y.append(label.numpy())

    return np.vstack(X), np.concatenate(y)

X768, y = load_and_process_split("public")
k = 4

set_seed()
km768 = KMeans(n_clusters=k, random_state=0).fit(X768)
public_clusters = km768.labels_

def transform(X_high):
    assert X_high.ndim == 2 and X_high.shape[1] >= 32, "Input must be 2D with >=32 dims."
    ### TODO: YOUR CODE HERE
    out = X_high[:, :32]
    ###

    assert out.ndim == 2 and out.shape[1] == 32, f"Embeddings must be shape (N,32), got {out.shape}."
    return out

def evaluate_clustering(ref_labels, n_clusters, n_runs=10):
    nmi_scores = []
    for seed in range(n_runs):
        set_seed(seed)
        X32 = transform(X768)
        km = KMeans(n_clusters=n_clusters, random_state=0).fit(X32)
        pred = km.labels_
        nmi = normalized_mutual_info_score(ref_labels, pred)
        nmi_scores.append(nmi)

    mean_nmi = np.mean(nmi_scores)
    std_nmi = np.std(nmi_scores)
    return mean_nmi, std_nmi

mean_nmi, std_nmi = evaluate_clustering(public_clusters, k)
print(f"\nPublic NMI (first-32 dims): {mean_nmi:.4f} ± {std_nmi:.4f}")