End-to-end, unsupervised clustering pipeline for large sets of text embeddings. It cleans the data with an ensemble of outlier detectors, chooses the number of clusters from a panel of four independent indices, and runs the final KMeans clustering. Built and validated on 45,895 texts embedded into a 3,072-dimensional space.
The algorithmic core lives in a plain, unit-tested module (cluster_pipeline.py); the Jupyter notebook is the exploratory walkthrough on top of it.
Three decisions drive the design. Each answers a failure mode that a naive fit() walks straight into.
- Clean before clustering. KMeans minimizes squared distance, so a handful of far-off points drag centroids toward themselves and inflate inertia. On real embedding data a few percent of items are junk (encoding errors, boilerplate, off-topic text). Removing them first is what separates a usable clustering from a distorted one.
- An ensemble of detectors, not one. Every outlier detector encodes a different, partial notion of "anomalous". A single detector is confidently wrong on the cases its assumption does not fit. Combining three independent views and requiring agreement turns three biased-but-cheap signals into one conservative decision.
- Four indices for k, not one magic number. There is no ground-truth k in unsupervised clustering, and every internal validity index has a known bias. Reporting four and reading their agreement (or disagreement) is more honest than trusting a single curve.
flowchart TD
A["Text embeddings<br/>n x 3072"] --> B{"Stage 1<br/>Ensemble outlier detection"}
B --> D1["KNN mean distance<br/>(global sparsity)"]
B --> D2["Local Outlier Factor<br/>(local density)"]
B --> D3["Isolation Forest<br/>(random-split isolation)"]
D1 --> V["Majority vote:<br/>drop point if >= 2 of 3 flag it"]
D2 --> V
D3 --> V
V --> C["Clean matrix"]
C --> K{"Stage 2<br/>Optimal k"}
K --> M1["Elbow (inertia)"]
K --> M2["Silhouette"]
K --> M3["Calinski-Harabasz"]
K --> M4["Davies-Bouldin"]
M1 --> MED["Consensus:<br/>median of the four"]
M2 --> MED
M3 --> MED
M4 --> MED
MED --> KM["Stage 3<br/>KMeans (k-means++)"]
KM --> T["t-SNE 2D<br/>(visualization only)"]
KM --> OUT["Labelled output CSV"]
Three detectors run independently, each capturing a different notion of "anomalous":
| Detector | What it catches | Signal |
|---|---|---|
| KNN mean distance (z-score > 1.0) | Points globally far from everything else | Mean distance to the 50 nearest neighbors, z-scored |
| Local Outlier Factor | Points in a locally sparse region even if not globally far | Local density vs the density of their neighbors |
| Isolation Forest | Points that are easy to isolate with random splits | Average path length in a random forest |
Aggregation rule. A point is removed only when at least 2 of 3 detectors flag it. This is a deliberately precision-oriented (conservative) vote: one detector misfiring is not enough to discard a sample, so the pipeline errs toward keeping data rather than silently deleting it.
| Method | Outliers flagged |
|---|---|
| KNN distance (z > 1.0) | 4,729 |
| Local Outlier Factor | 4,590 |
| Isolation Forest | 4,590 |
| Ensemble (>= 2 votes) | 2,591 |
Result on the reference data: 45,895 -> 43,304 clean samples (5.6% removed). The ensemble removes far fewer points than any single detector, which is the intended effect of requiring consensus.
Four standard internal validity indices evaluate every k from 2 to 80:
| Index | Optimizes | Bias |
|---|---|---|
| Elbow (inertia) | Diminishing returns of added clusters | Toward few, broad clusters |
| Silhouette | Cohesion vs separation (max) | Toward few, broad clusters |
| Calinski-Harabasz | Between/within variance ratio (max) | Toward many, fine clusters |
| Davies-Bouldin | Avg similarity to nearest cluster (min) | Toward many, fine clusters |
| Index | Optimal k |
|---|---|
| Elbow | 2 |
| Silhouette | 2 |
| Calinski-Harabasz | 57 |
| Davies-Bouldin | 57 |
| Consensus (median of the four) | 29 |
Consensus rule. The consensus k is the median of the four, int(np.median([2, 2, 57, 57])) = 29: a robust aggregate that no single divergent index can dominate.
Reading the disagreement. The split is the finding, not a bug. Compactness metrics (elbow, silhouette) see about 2 broad groups; separation metrics (Calinski-Harabasz, Davies-Bouldin) see about 57 fine ones. That gap is the signature of hierarchical structure: a few large semantic groups with finer structure nested inside. The consensus 29 is a neutral midpoint to start from, not a validated optimum; in production k should be driven by the downstream task, and the pipeline exposes a k_override hook for exactly that.
KMeans with k-means++ initialization runs in the original 3,072-dimensional space (t-SNE is used only to project the result to 2D for a sanity-check plot, never for the clustering itself). Output is the original rows plus a Cluster label, written to CSV.
The senior-level choices, in one place:
- Consensus over single-shot at both risky steps: 2-of-3 vote for cleaning, median-of-4 for k. Both trade a little recall for robustness against any one method's blind spot.
- Conservative cleaning. The vote is tuned to keep data (>= 2 votes to drop), because deleting a real sample is more expensive than keeping a borderline one in an exploratory pipeline.
- Cluster in full dimension, visualize in 2D. t-SNE distances are not metric-faithful, so it informs the eye but never the algorithm.
- Edge cases handled explicitly. Degenerate input (all-identical points, zero variance) returns "no outliers" instead of dividing by zero;
random_stateis threaded through every estimator so runs are reproducible. - Logic is a library, not a notebook. The reusable core is pure functions with type hints and no I/O, so it is importable, testable and CI-gated. The notebook consumes it for exploration.
cluster_pipeline.py # Reusable core: detection, k-selection, clustering (pure, tested)
cluster_optimization.ipynb # Exploratory notebook: load, run, plot, narrate
tests/test_cluster_pipeline.py # Behavioural tests on synthetic data with known ground truth
.github/workflows/ci.yml # CI: ruff lint + pytest on Python 3.10 / 3.11 / 3.12
cluster_pipeline.py exposes the pipeline as composable functions plus a one-call run_pipeline() orchestrator:
from cluster_pipeline import run_pipeline
result = run_pipeline(matrix, max_clusters=80) # metric-driven k
result = run_pipeline(matrix, k_override=12) # task-driven k
result.inlier_mask # bool mask of kept rows
result.k_selection # per-method optimal k + consensus
result.labels # final cluster assignments- Python 3.10+
- scikit-learn: KMeans, t-SNE, LOF, Isolation Forest, NearestNeighbors, silhouette / Calinski-Harabasz / Davies-Bouldin scores
- NumPy / pandas: numerics and data handling
- matplotlib / seaborn: visualization (notebook)
- pytest / ruff: tests and linting, run in CI
pip install -r requirements.txtRun the reusable pipeline directly:
import numpy as np
from cluster_pipeline import run_pipeline
matrix = np.vstack(df["embedding"].values) # n x d float array
result = run_pipeline(matrix, max_clusters=80)
df["Cluster"] = -1
df.loc[result.inlier_mask, "Cluster"] = result.labelsOr open cluster_optimization.ipynb in Jupyter / Colab and run all cells. The notebook expects a CSV with columns _id, text, n_tokens, embedding (a string-encoded float array).
The test suite validates behaviour, not just imports. On synthetic data with a known ground truth it checks that:
- planted far-off outliers are flagged by each detector and removed by the ensemble, while the genuine blob (>= 85%) survives;
- the ensemble is never more aggressive than a single detector (consensus property);
- degenerate zero-variance input does not divide by zero;
- on three well-separated Gaussian blobs, silhouette, Calinski-Harabasz and Davies-Bouldin all recover k = 3, and each blob lands in a single cluster;
- outputs are shape-correct and reproducible under a fixed seed, end to end.
pip install -r requirements.txt pytest
pytest -qCI runs ruff and the full suite on Python 3.10, 3.11 and 3.12 on every push and pull request.
MIT. See LICENSE.