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