src/clustering.py

import numpy as np
import faiss

class ClusterManager:
    def __init__(self, embeddings, n_clusters=100):
        self.n_clusters = n_clusters
        self.dim = embeddings.shape[1]
        
        self.quantizer = faiss.IndexFlatIP(self.dim)
        self.index = faiss.IndexIVFFlat(self.quantizer, self.dim, self.n_clusters, faiss.METRIC_INNER_PRODUCT)
        self.index.train(embeddings)
        self.index.add(embeddings)
        
        self.centroids = np.vstack([self.index.quantizer.reconstruct(i) for i in range(self.n_clusters)]).astype("float32")
        _, doc_cluster_assign = self.index.quantizer.search(embeddings, 1)
        self.doc_cluster_ids = doc_cluster_assign.flatten()
        
        self.cluster_to_docs = {i: [] for i in range(self.n_clusters)}
        for doc_i, c in enumerate(self.doc_cluster_ids):
            self.cluster_to_docs[c].append(doc_i)
            
        self.m_neighbors = min(128, self.n_clusters - 1)
        self.centroid_sims = np.dot(self.centroids, self.centroids.T)
        self.cluster_neighbors = {}
        for c in range(self.n_clusters):
            sims = self.centroid_sims[c].copy()
            sims[c] = -1.0 
            self.cluster_neighbors[c] = np.argsort(sims)[::-1][:self.m_neighbors].tolist()