Coverage for src/lilbee/retrieval/clustering_embedding/helpers.py: 100%
168 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Numeric and tokenization helpers for the embedding clusterer.
3Mutual-kNN is hub-robust: a pathological hub ends up in many one-way
4neighborhoods but can reciprocate at most ``k`` of them, so hub-driven
5bridging across topics is broken at the graph-construction step without
6any post-hoc similarity rescaling.
8The similarity kernel is blocked in row chunks of ``_BLOCK_SIZE`` to keep
9peak memory bounded regardless of corpus size, so it scales comfortably
10to tens of thousands of chunks on a laptop.
11"""
13from __future__ import annotations
15import math
16from collections import Counter
17from collections.abc import Mapping
19import numpy as np
21from lilbee.data.store import Store
22from lilbee.retrieval.clustering import SourceCluster
23from lilbee.retrieval.clustering_embedding.types import ClusterChunk
25# Block size for the similarity kernel. With N=10000 and D=768 this caps
26# peak float32 memory at block * N * 4 bytes ~= 40 MB.
27_BLOCK_SIZE = 1024
29# Rows boxed into Python objects per batch when scanning the chunks table.
30# Caps the transient working set at ~25 MB (1024 rows x 768 dims), whatever
31# the corpus size.
32_SCAN_BATCH_ROWS = 1024
33# The chunk columns a cluster scan reads.
34_SCAN_COLUMNS = ["source", "chunk", "chunk_index", "vector"]
36# Label Propagation hard iteration cap. Convergence is typically reached
37# in well under 10 passes on real corpora.
38_MAX_LPA_ITERATIONS = 30
40# Minimum non-zero L2 norm for a row vector to be kept.
41_MIN_VECTOR_NORM = 1e-12
43# A source joins a chunk community on at least
44# `min(_MIN_SOURCE_CHUNKS, ceil(total * _MIN_SOURCE_FRACTION))` of its chunks.
45# min() is the lenient cutoff: it caps the requirement at _MIN_SOURCE_CHUNKS so
46# a long document needs a foothold, not a proportional chunk count.
47_MIN_SOURCE_CHUNKS = 3
48_MIN_SOURCE_FRACTION = 0.2
50# TF-IDF labeling knobs.
51_LABEL_TOP_TERMS = 3
52# Chunks with fewer tokens than this are down-weighted when accumulating
53# term frequency so short boilerplate (headings, captions) cannot dominate
54# a cluster label. 20 tokens roughly matches the token count of a section
55# heading or a two-sentence summary.
56_SHORT_CHUNK_TOKEN_CAP = 20
58# kNN auto-scaling bounds. Formula: clamp(round(log2(N)+2), _MIN_K, _MAX_K).
59_MIN_K = 5
60_MAX_K = 20
63def auto_k(n: int) -> int:
64 """Pick a neighborhood size from corpus size via ``clamp(log2(N)+2)``."""
65 if n <= 1:
66 return _MIN_K
67 raw = round(math.log2(max(n, 4)) + 2)
68 return max(_MIN_K, min(_MAX_K, raw))
71def _parse_chunk_row(
72 row: Mapping[str, object],
73) -> tuple[ClusterChunk, list[float] | tuple[float, ...]] | None:
74 """Extract a chunk record + vector from a raw Arrow row, or None on invalid."""
75 vector = row.get("vector")
76 if not isinstance(vector, (list, tuple)):
77 return None
78 source = row.get("source")
79 if not isinstance(source, str):
80 return None
81 raw_text = row.get("chunk")
82 chunk_text = raw_text if isinstance(raw_text, str) else ""
83 raw_index = row.get("chunk_index")
84 chunk_index = raw_index if isinstance(raw_index, int) else 0
85 # tokens derive from text in ClusterChunk.__post_init__.
86 record = ClusterChunk(source=source, chunk_index=chunk_index, text=chunk_text)
87 return record, vector
90def _load_chunk_records(
91 store: Store,
92) -> tuple[list[ClusterChunk], np.ndarray]:
93 """Scan the chunks table once and return records plus a float32 matrix.
95 Rows with an unparseable vector are skipped. Records are sorted by
96 ``(source, chunk_index)`` so downstream cluster IDs are stable
97 regardless of LanceDB's row return order. Records are tokenized once
98 here so TF-IDF labeling does not re-tokenize.
100 The scan consumes the table in bounded row batches (``_SCAN_BATCH_ROWS``):
101 each batch's rows are parsed and its vectors immediately compacted into a
102 float32 block, so the boxed Python working set (row dicts plus vector
103 float lists) is one batch, not the whole table. Only the compact matrix
104 and the text records scale with corpus size.
105 """
106 records: list[ClusterChunk] = []
107 blocks: list[np.ndarray] = []
108 for rows in store.chunk_batches(_SCAN_COLUMNS, _SCAN_BATCH_ROWS):
109 batch_vectors: list[list[float] | tuple[float, ...]] = []
110 for row in rows:
111 pair = _parse_chunk_row(row)
112 if pair is None:
113 continue
114 record, vector = pair
115 records.append(record)
116 batch_vectors.append(vector)
117 if batch_vectors:
118 blocks.append(np.asarray(batch_vectors, dtype=np.float32))
119 if not records:
120 return [], np.zeros((0, 0), dtype=np.float32)
122 matrix = np.vstack(blocks)
123 order = sorted(range(len(records)), key=lambda i: (records[i].source, records[i].chunk_index))
124 return [records[i] for i in order], matrix[order]
127def normalize_rows(matrix: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
128 """Return (normalized_matrix, keep_mask). Zero-norm rows are dropped."""
129 if matrix.size == 0:
130 return matrix, np.zeros(0, dtype=bool)
131 norms = np.linalg.norm(matrix, axis=1)
132 keep = norms > _MIN_VECTOR_NORM
133 if not keep.all():
134 matrix = matrix[keep]
135 norms = norms[keep]
136 return matrix / norms[:, None], keep
139def mutual_knn(matrix: np.ndarray, k: int) -> dict[int, set[int]]:
140 """Build a mutual k-nearest-neighbors graph over L2-normalized rows.
142 Computes similarity in row blocks so peak memory stays bounded.
143 Self-similarity is masked so each row's neighbors exclude itself.
144 Mutuality is enforced by keeping only edges ``(i, j)`` where
145 ``j`` is in row ``i``'s top-k AND ``i`` is in row ``j``'s top-k;
146 this single rule breaks hub-driven bridging without any extra
147 similarity rescaling.
148 """
149 n = matrix.shape[0]
150 if n == 0 or k <= 0:
151 return {}
153 effective_k = min(k, n - 1)
154 if effective_k <= 0:
155 return {i: set() for i in range(n)}
157 top_neighbors: list[set[int]] = [set() for _ in range(n)]
159 for start in range(0, n, _BLOCK_SIZE):
160 stop = min(start + _BLOCK_SIZE, n)
161 sim_block = matrix[start:stop] @ matrix.T # (block, n)
162 # Mask self-similarity so each row's own index is never returned.
163 # For the tail block where stop-start < _BLOCK_SIZE the fancy-index
164 # pairing still lines up because both sides are length stop-start.
165 block_rows = np.arange(stop - start)
166 sim_block[block_rows, np.arange(start, stop)] = -math.inf
167 # Partition by largest similarities without allocating a negated
168 # copy of the block: pass a negative kth to select the tail.
169 neighbor_idx = np.argpartition(sim_block, -effective_k, axis=1)[:, -effective_k:]
170 for local_row, global_row in enumerate(range(start, stop)):
171 top_neighbors[global_row] = set(neighbor_idx[local_row].tolist())
173 mutual: dict[int, set[int]] = {i: set() for i in range(n)}
174 for i in range(n):
175 for j in top_neighbors[i]:
176 if i in top_neighbors[j]:
177 mutual[i].add(j)
178 return mutual
181def label_propagation(
182 adjacency: dict[int, set[int]],
183 order: list[int],
184) -> list[int]:
185 """Async Label Propagation with deterministic min-label tie-breaking.
187 Each node adopts the most common label among its neighbors. Ties are
188 broken by smallest label id so the outcome is reproducible across
189 runs on the same corpus. The caller is responsible for passing a
190 complete ``adjacency`` (one entry per node 0..n-1) and an ``order``
191 that covers every node: ``get_clusters`` guarantees both.
192 """
193 n = len(adjacency)
194 labels = list(range(n))
195 for _ in range(_MAX_LPA_ITERATIONS):
196 changed = False
197 for node in order:
198 neighbors = adjacency.get(node)
199 if not neighbors:
200 continue
201 counts: Counter[int] = Counter(labels[j] for j in neighbors)
202 top_count = max(counts.values())
203 best = min(label for label, count in counts.items() if count == top_count)
204 if labels[node] != best:
205 labels[node] = best
206 changed = True
207 if not changed:
208 break
209 return labels
212def communities_by_label(labels: list[int]) -> dict[int, list[int]]:
213 """Group node indices by their final community label."""
214 communities: dict[int, list[int]] = {}
215 for node, label in enumerate(labels):
216 communities.setdefault(label, []).append(node)
217 return communities
220def _source_totals(records: list[ClusterChunk]) -> dict[str, int]:
221 """Return the total chunk count per source across the whole corpus."""
222 totals: dict[str, int] = {}
223 for record in records:
224 totals[record.source] = totals.get(record.source, 0) + 1
225 return totals
228def _filter_sources(
229 member_indices: list[int],
230 records: list[ClusterChunk],
231 source_totals: dict[str, int],
232) -> frozenset[str]:
233 """Apply the source-membership threshold to a community's members."""
234 per_source: dict[str, int] = {}
235 for idx in member_indices:
236 source = records[idx].source
237 per_source[source] = per_source.get(source, 0) + 1
238 kept: set[str] = set()
239 for source, count in per_source.items():
240 total = source_totals.get(source, count)
241 fractional_cutoff = math.ceil(total * _MIN_SOURCE_FRACTION)
242 cutoff = min(_MIN_SOURCE_CHUNKS, fractional_cutoff)
243 if count >= cutoff:
244 kept.add(source)
245 return frozenset(kept)
248def _corpus_document_frequency(records: list[ClusterChunk]) -> dict[str, int]:
249 """Compute document frequency (chunk count containing term) for every term."""
250 df: dict[str, int] = {}
251 for record in records:
252 for term in set(record.tokens):
253 df[term] = df.get(term, 0) + 1
254 return df
257def _label_community(
258 member_indices: list[int],
259 records: list[ClusterChunk],
260 df: dict[str, int],
261 total_chunks: int,
262 fallback: str,
263) -> str:
264 """Pick a topic label for a community using sublinear TF-IDF scoring."""
265 tf: dict[str, float] = {}
266 for idx in member_indices:
267 tokens = records[idx].tokens
268 if not tokens:
269 continue
270 weight = min(1.0, len(tokens) / _SHORT_CHUNK_TOKEN_CAP)
271 counts: Counter[str] = Counter(tokens)
272 for term, count in counts.items():
273 tf[term] = tf.get(term, 0.0) + weight * (1.0 + math.log(count))
275 scored: list[tuple[float, str]] = []
276 for term, term_tf in tf.items():
277 # Standard ``log(N / (1 + df))`` smoothing: the +1 keeps the
278 # denominator non-zero for new terms and damps the score of
279 # terms that appear in every chunk (where idf goes negative
280 # and the term is filtered out entirely).
281 idf = math.log(total_chunks / (1 + df.get(term, 0)))
282 if idf <= 0:
283 continue
284 scored.append((term_tf * idf, term))
285 if not scored:
286 return fallback
288 scored.sort(key=lambda pair: (-pair[0], pair[1]))
289 return " ".join(term for _, term in scored[:_LABEL_TOP_TERMS])
292def _build_clusters(
293 communities: dict[int, list[int]],
294 records: list[ClusterChunk],
295 source_totals: dict[str, int],
296 df: dict[str, int],
297 min_sources: int,
298) -> tuple[list[SourceCluster], int]:
299 """Turn raw chunk communities into published source clusters.
301 Returns ``(clusters, noise_chunk_count)`` where ``noise_chunk_count``
302 is the number of chunks whose community failed the source filter.
303 """
304 ordered = sorted(communities.items(), key=lambda pair: (-len(pair[1]), pair[0]))
305 total_chunks = len(records)
306 clusters: list[SourceCluster] = []
307 noise = 0
308 for idx, (_, members) in enumerate(ordered):
309 kept_sources = _filter_sources(members, records, source_totals)
310 if len(kept_sources) < min_sources:
311 noise += len(members)
312 continue
313 cluster_id = f"embedding-{idx}"
314 label = _label_community(members, records, df, total_chunks, fallback=cluster_id)
315 clusters.append(
316 SourceCluster(
317 cluster_id=cluster_id,
318 label=label,
319 sources=kept_sources,
320 )
321 )
322 return clusters, noise