From 8a19c6aa5a45af3891c9fff10441eaef680b5fc4 Mon Sep 17 00:00:00 2001 From: Mattie Schraeder Date: Mon, 6 Jul 2026 08:09:35 -0500 Subject: [PATCH] Make RAPTOR GMM robust on small reduced clusters (#16632) --- rag/advanced_rag/knowlege_compile/raptor.py | 8 ++++++-- test/unit_test/rag/test_raptor_psi_tree_builder.py | 2 +- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/rag/advanced_rag/knowlege_compile/raptor.py b/rag/advanced_rag/knowlege_compile/raptor.py index 97c934d9d..cda5779e1 100644 --- a/rag/advanced_rag/knowlege_compile/raptor.py +++ b/rag/advanced_rag/knowlege_compile/raptor.py @@ -43,6 +43,10 @@ from rag.utils.raptor_utils import ( SUPPORTED_TREE_BUILDERS, ) +# Regularization added to GMM covariance diagonals; keeps components +# from collapsing on singleton/near-identical reduced points. +_GMM_REG_COVAR = 1e-4 + @dataclass class _PsiTreeNode: @@ -254,7 +258,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval: for n in n_clusters: self._check_task_canceled(task_id, "get optimal clusters") - gm = GaussianMixture(n_components=n, random_state=random_state) + gm = GaussianMixture(n_components=n, random_state=random_state, covariance_type="diag", reg_covar=_GMM_REG_COVAR) gm.fit(embeddings) bics.append(gm.bic(embeddings)) optimal_clusters = n_clusters[np.argmin(bics)] @@ -344,7 +348,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval: if n_clusters <= 1: labels = [0 for _ in range(len(reduced_embeddings))] else: - gm = GaussianMixture(n_components=n_clusters, random_state=random_state) + gm = GaussianMixture(n_components=n_clusters, random_state=random_state, covariance_type="diag", reg_covar=_GMM_REG_COVAR) gm.fit(reduced_embeddings) probs = gm.predict_proba(reduced_embeddings) labels = [] diff --git a/test/unit_test/rag/test_raptor_psi_tree_builder.py b/test/unit_test/rag/test_raptor_psi_tree_builder.py index aa4a16101..de8465b3e 100644 --- a/test/unit_test/rag/test_raptor_psi_tree_builder.py +++ b/test/unit_test/rag/test_raptor_psi_tree_builder.py @@ -245,7 +245,7 @@ def test_get_optimal_clusters_evaluates_upper_bound_candidate(monkeypatch, rapto evaluated = [] class RecordingGaussianMixture: - def __init__(self, n_components, random_state=None): + def __init__(self, n_components, random_state=None, **kwargs): self.n_components = n_components evaluated.append(n_components)