zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-14 08:58:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

Make RAPTOR GMM robust on small reduced clusters (#16632)

Browse Source
This commit is contained in:
Mattie Schraeder
2026-07-06 08:09:35 -05:00
committed by GitHub
parent 85b565244d
commit 8a19c6aa5a
2 changed files with 7 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
rag/advanced_rag/knowlege_compile/raptor.py
View File

@@ -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 = []