hcr — Hierarchical Clustering Refinement

hcr is a pip-installable Python package for hierarchical KMeans clustering with post-hoc refinement.

Bring your own embeddings. hcr operates on any L2-normalised float32 numpy array — no text loading, no model dependencies baked in.

---

Install

pip install hcr

For UMAP visualisations:

pip install "hcr[visualization]"

---

Quickstart

import numpy as np
from hcr import HierarchicalClusteringRefinement

# Your embeddings: shape (N, D), float32, L2-normalised rows.
# Produce them however you like (sentence-transformers, OpenAI, E5, etc.)
X = np.load("embeddings.npy").astype(np.float32)

# k is auto-estimated from the data when not specified
hcr = HierarchicalClusteringRefinement(seed=42)
result = hcr.fit(X)

print(result.labels)      # np.ndarray, shape (N,) — cluster id per row, -1 = outlier
print(result.resolved_k)  # {"first_level_k": 14, "second_level_k": 14, "leaf_k": 14, "auto_estimated": True}
print(result.summary)     # dict: counts of merges, removals, reassignments
print(result.tree)        # pd.DataFrame: the hierarchical tree structure

Explicit k values

# All three levels specified — used as-is
hcr = HierarchicalClusteringRefinement(first_level_k=20, second_level_k=15, leaf_k=10, seed=42)

# Only first two — leaf_k inherits from second_level_k
hcr = HierarchicalClusteringRefinement(first_level_k=20, second_level_k=15, seed=42)

# Only first — second_level_k and leaf_k both inherit from first_level_k
hcr = HierarchicalClusteringRefinement(first_level_k=20, seed=42)

With text for TF-IDF cluster reports

texts = ["some sentence", "another sentence", ...]  # len(texts) == N

result = hcr.fit(X, texts=texts)
print(result.report)    # pd.DataFrame: per-cluster stats + top TF-IDF terms + examples

---

Adaptive k selection

All three k parameters (first_level_k, second_level_k, leaf_k) default to None. The values used are resolved inside fit() according to this priority:

Provided	first_level_k	second_level_k	leaf_k
All three	as given	as given	as given
first + second	as given	as given	= second_level_k
first only	as given	= first_level_k	= first_level_k
none	estimated	estimated	estimated

When none are given, estimate_k_from_similarity_distribution(X) is called. It samples random pairs from X, fits a Gaussian (μ, σ) to the cosine similarity distribution, and returns:

k = clamp(int(1 / (1 - μ + σ)), min=2, max=50)

A high mean similarity (dense, similar embeddings) produces a larger k; a wider spread (σ) inflates k further. The resolved values are always stored on the result for inspection:

print(result.resolved_k)
# {"first_level_k": 14, "second_level_k": 14, "leaf_k": 14, "auto_estimated": True}

Control sampling cost with k_estimation_sample_pairs (default 50_000).

---

ClusteringResult fields

Field	Type	Description
labels	np.ndarray (N,)	Refined cluster label per row; -1 = outlier
raw_labels	np.ndarray (N,)	Pre-refinement leaf labels
resolved_k	dict	k values used + "auto_estimated" bool
summary	dict	Counts of removes, merges, reassignments, outliers
report	pd.DataFrame | None	Per-cluster stats + TF-IDF terms; None if no texts given
tree	pd.DataFrame	Full hierarchical tree (one row per node)
assignments	pd.DataFrame	Row-level path/depth/node from the clustering step

---

CLI

hcr --embeddings embeddings.npy \
    --data texts.csv --text-col text \
    --first-level-k 20 --second-level-k 20 \
    --output-dir runs_h

Omit all k flags to use auto-estimation:

hcr --embeddings embeddings.npy --output-dir runs_h

Key flags:

Flag	Default	Description
--first-level-k	auto	Clusters at depth 1
--second-level-k	auto	Clusters at depth 2
--leaf-k	auto	k for adaptive leaf splits
--k-estimation-sample-pairs	50000	Pairs sampled for auto k estimation
--sim-threshold	0.80	Cosine similarity target for leaf quality
--refine-merge-threshold	0.90	Similarity floor to attempt cluster merging
--seed	42	Global random seed

Run hcr --help for all options.

---

Key parameters

Parameter	Default	Description
first_level_k	None (auto)	Clusters at depth 1
second_level_k	None (auto)	Clusters at depth 2
leaf_k	None (auto)	k for adaptive splits at leaf nodes
k_estimation_sample_pairs	50_000	Random pairs sampled for k auto-estimation
sim_threshold	0.80	Cosine similarity target for leaf quality
alpha	0.6	Centring strength at each recursive level
refine_remove_abs_threshold	0.45	Hard similarity floor for member removal
refine_merge_threshold	0.90	Minimum rep-similarity to attempt merging
refine_reassign_threshold	0.60	Similarity needed to reassign a removed item
seed	42	Global random seed (used for k estimation too)

---

Package structure

hcr/
├── __init__.py        HierarchicalClusteringRefinement, ClusteringResult
├── clustering.py      hierarchical_kmeans, estimate_k_from_similarity_distribution, …
├── refinement.py      remove_bad_members_B_rule, greedy_iterative_merge, …
├── reporting.py       cluster_report, build_leaf_row_table, …
├── sampling.py        kcenter_greedy_indices
├── preprocessing.py   center_then_l2, l2_normalize_rows, mask_numbers, …
├── visualization.py   reduce_and_plot, plot_all_node_tsnes, …
└── cli.py             Config, parse_args, main
tests/
└── test_smoke.py      pytest smoke tests (no external data, no torch)

---

Background

Hierarchical Clustering Refinement (HCR) is a post-clustering framework designed to improve the quality of hierarchical clustering results in large-scale intent discovery, semantic grouping, and embedding-based clustering tasks.

Traditional clustering methods such as K-Means or Hierarchical K-Means often suffer from structural limitations when applied to high-dimensional embedding spaces:

• forced assignment of samples to incorrect clusters
• cluster fragmentation caused by recursive splits
• semantic overlap between clusters
• noisy samples degrading cluster representations

HCR introduces a refinement stage applied after hierarchical clustering to repair these issues and produce more coherent clusters.

---

Motivation

When clustering large embedding datasets (e.g., customer messages, search queries, support tickets), hierarchical K-Means provides scalability and interpretability, but it often produces:

• over-segmented clusters
• semantically duplicated clusters
• clusters contaminated with outliers
• samples assigned to the wrong branch during early splits

HCR addresses these problems by introducing a structured cluster repair stage that operates on the final leaf clusters.

---

Method Overview

Hierarchical Clustering Refinement (HCR) is designed as a post-clustering refinement framework that operates on the output of hierarchical clustering algorithms such as Hierarchical K-Means.

The overall workflow consists of two main phases:

1. Hierarchical Clustering
   The dataset is first clustered using a hierarchical strategy (e.g., recursive K-Means). This stage produces the initial set of leaf clusters.

2. Cluster Refinement (HCR)
   HCR is then applied to the resulting clusters to repair common clustering issues such as outliers, fragmented clusters, and incorrect assignments.

---

Pipeline

The full processing pipeline can be summarized as follows:

flowchart TD
    A[Embedding Generation] --> B[Hierarchical K-Means Clustering]
    B --> C[Leaf Clusters]
    C --> D[HCR Refinement Stage]

    D --> D1[Outlier Removal]
    D1 --> D2[Greedy Cluster Merge]
    D2 --> D3[Sample Reassignment]

    D3 --> E[Improved Clusters]

Loading

---

Refinement Stages

The refinement process consists of three core steps.

1. Outlier Removal

Cluster members that are significantly dissimilar from their cluster representation are removed.

Instead of using a fixed threshold alone, HCR uses a hybrid rule combining absolute and relative thresholds:

remove if similarity < max(
    absolute_threshold,
    cluster_p10 - margin
)

Where:

• cluster_p10 is the 10th percentile similarity within the cluster
• margin is a tolerance parameter

This allows clusters with different densities to be treated adaptively.

---

2. Greedy Cluster Merge

Clusters that are extremely similar are iteratively merged using a greedy similarity-based merge strategy.

The algorithm repeatedly:

1. computes similarity between cluster representatives
2. finds the most similar cluster pair
3. merges them if similarity exceeds a threshold
4. recomputes the merged cluster representation

This process continues until no cluster pair exceeds the merge threshold.

This approach prevents chain merges commonly observed in graph-based merging strategies.

---

3. Sample Reassignment

Samples removed during the outlier stage are evaluated again.

For each removed sample:

assign to cluster if max_similarity ≥ reassign_threshold
otherwise mark as outlier

This allows incorrectly assigned samples to recover and join the correct cluster.

---

Cluster Representation

Cluster representatives are computed using the mean embedding of cluster members, followed by normalization.

cluster_rep = normalize(mean(cluster_member_embeddings))

Using the actual cluster mean instead of K-Means centroids ensures that the representative vector accurately reflects the final cluster structure.

---

Key Features

• adaptive k selection — auto-estimated from the similarity distribution when not specified
• hierarchical clustering compatible
• scalable to large embedding datasets
• adaptive outlier detection (hybrid absolute + relative threshold)
• greedy cluster consolidation with internal quality guards
• semantic reassignment of removed samples
• embedding-model agnostic

---

Use Cases

HCR is particularly useful for:

• intent discovery
• customer support message clustering
• search query grouping
• topic discovery in large text datasets
• semantic dataset exploration

---

Example Workflow

1. Generate sentence embeddings
2. Run hierarchical K-Means clustering
3. Extract final leaf clusters
4. Apply HCR refinement
5. Analyze improved cluster structure

---

License

See the LICENSE file for full license text.