cat-stack 2.4.0__tar.gz → 2.5.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (45) hide show
  1. {cat_stack-2.4.0 → cat_stack-2.5.1}/PKG-INFO +9 -5
  2. {cat_stack-2.4.0 → cat_stack-2.5.1}/README.md +6 -2
  3. {cat_stack-2.4.0 → cat_stack-2.5.1}/pyproject.toml +2 -2
  4. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/__about__.py +1 -1
  5. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/collapse_themes.py +241 -11
  6. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/explore.py +5 -3
  7. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/extract.py +109 -6
  8. {cat_stack-2.4.0 → cat_stack-2.5.1}/.gitignore +0 -0
  9. {cat_stack-2.4.0 → cat_stack-2.5.1}/LICENSE +0 -0
  10. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/cat_stack/__init__.py +0 -0
  11. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/__init__.py +0 -0
  12. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/_batch.py +0 -0
  13. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/_category_analysis.py +0 -0
  14. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/_chunked.py +0 -0
  15. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/_embeddings.py +0 -0
  16. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/_formatter.py +0 -0
  17. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/_pilot_test.py +0 -0
  18. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/_prompts.py +0 -0
  19. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/_providers.py +0 -0
  20. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/_review_ui.py +0 -0
  21. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/_tiebreaker.py +0 -0
  22. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/_utils.py +0 -0
  23. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/_web_fetch.py +0 -0
  24. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/_wrapper_helpers.py +0 -0
  25. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/calls/CoVe.py +0 -0
  26. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/calls/__init__.py +0 -0
  27. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/calls/image_CoVe.py +0 -0
  28. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/calls/image_stepback.py +0 -0
  29. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/calls/pdf_CoVe.py +0 -0
  30. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/calls/pdf_stepback.py +0 -0
  31. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/calls/stepback.py +0 -0
  32. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/calls/top_n.py +0 -0
  33. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/classify.py +0 -0
  34. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/image_functions.py +0 -0
  35. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/images/circle.png +0 -0
  36. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/images/cube.png +0 -0
  37. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/images/diamond.png +0 -0
  38. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/images/overlapping_pentagons.png +0 -0
  39. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/images/rectangles.png +0 -0
  40. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/model_reference_list.py +0 -0
  41. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/pdf_functions.py +0 -0
  42. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/prompt_tune.py +0 -0
  43. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/summarize.py +0 -0
  44. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/text_functions.py +0 -0
  45. {cat_stack-2.4.0 → cat_stack-2.5.1}/src/catstack/text_functions_ensemble.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: cat-stack
3
- Version: 2.4.0
3
+ Version: 2.5.1
4
4
  Summary: Domain-agnostic text, image, PDF, and DOCX classification engine powered by LLMs
5
5
  Project-URL: Documentation, https://github.com/chrissoria/cat-stack#readme
6
6
  Project-URL: Issues, https://github.com/chrissoria/cat-stack/issues
@@ -24,9 +24,9 @@ Requires-Dist: pandas
24
24
  Requires-Dist: requests
25
25
  Requires-Dist: tqdm
26
26
  Provides-Extra: agent
27
- Requires-Dist: cat-claws[claude]>=0.3.0; extra == 'agent'
27
+ Requires-Dist: cat-claws[claude]>=0.3.1; extra == 'agent'
28
28
  Provides-Extra: codex-agent
29
- Requires-Dist: cat-claws[codex]>=0.3.0; extra == 'codex-agent'
29
+ Requires-Dist: cat-claws[codex]>=0.3.1; extra == 'codex-agent'
30
30
  Provides-Extra: docx
31
31
  Requires-Dist: python-docx>=1.0.0; extra == 'docx'
32
32
  Provides-Extra: embeddings
@@ -138,7 +138,7 @@ cat.classify(
138
138
  ```
139
139
 
140
140
  ### `extract()`
141
- Discover categories from a corpus using LLM-driven exploration.
141
+ Discover categories from a corpus using LLM-driven exploration. Since v2.5.0, text consolidation runs the full explore → `collapse_themes()` pipeline (`engine="collapse"`, the default): the entire raw label inventory reaches the semantic merge — Jaro-Winkler dedup, embedding pre-merge, quality-controlled LLM passes, then a count-guided reduction to at most `max_categories`. Pass `engine="legacy"` to reproduce pre-2.5 runs (single merge call over a truncated inventory). `collapse_kwargs` forwards options to `collapse_themes()`; `max_workers` parallelizes both extraction and consolidation.
142
142
 
143
143
  ```python
144
144
  cat.extract(
@@ -197,7 +197,7 @@ cat.collapse_themes(
197
197
  | Parameter | Default | Description |
198
198
  | --- | --- | --- |
199
199
  | `input_data` | — | List of category labels, or a frequency `Series`/`dict` (`label -> count`). |
200
- | `api_key` | `None` | API key for the LLM provider (required). |
200
+ | `api_key` | `None` | API key for the LLM provider. Not required for subscription/CLI backends (`claude-code`, `claude-agent`, `codex-agent`) or `ollama`. |
201
201
  | `description` | `""` | The survey question or context, used in the merge prompt. |
202
202
  | `passes` | `1` | Number of merge iterations, or `"auto"` to iterate until the embedding-quality benchmark peaks. |
203
203
  | `max_passes` | `10` | Cap on iterations when `passes="auto"`. |
@@ -207,6 +207,8 @@ cat.collapse_themes(
207
207
  | `embedding_merge_threshold` | `0.92` | Cosine similarity at/above which labels are merged in the pre-LLM embedding step. `None`/`>=1.0` disables it. |
208
208
  | `shuffle` | `True` | Randomize order each pass so batch composition varies (improves convergence stability). |
209
209
  | `final_consolidation` | `0.82` | Cosine threshold for one greedy global embedding re-merge after all passes, collapsing cross-batch duplicates. Conservative by design (errs toward keeping categories). `False`/`None` skips it. |
210
+ | `top_n` | `None` | If set, a final global LLM call consolidates the surviving list into at most N categories, guided by each label's generation count (frequent themes favored; overlapping labels merged rather than dropped). Guaranteed `<= top_n` via truncation plus a deterministic top-N-by-count fallback. |
211
+ | `prune` | `False` | `True` = drop conceptual duplicates keeping one representative verbatim; never renames or merges merely-related labels. |
210
212
  | `user_model` | `"gpt-4o"` | Model for the merge phase. Use a capable model — small models can degenerate. |
211
213
  | `model_source` | `"auto"` | Provider for `user_model` (`"auto"`, `"openai"`, `"huggingface"`, …). |
212
214
  | `unique_model` | `None` | If set, run an initial extract-unique thinning phase on this (typically cheaper) model before the merge phase. `None` skips the phase (backward compatible). |
@@ -237,6 +239,8 @@ OpenAI, Anthropic, Google (Gemini), Mistral, Perplexity, xAI (Grok), HuggingFace
237
239
 
238
240
  All providers use the same `(model_name, provider, api_key)` tuple format. Provider is auto-detected from model name if omitted.
239
241
 
242
+ **Subscription backends (no API key).** Three `model_source` values authenticate through a chat subscription instead of a metered key — leave `api_key` unset: `"claude-agent"` (Claude subscription via the Agent SDK; `pip install "cat-stack[agent]"`), `"claude-code"` (the Claude Code CLI, if installed — no extra needed), and `"codex-agent"` (ChatGPT subscription; `pip install "cat-stack[codex-agent]"`). Classification, extraction, exploration, and summarization all route through them.
243
+
240
244
  ## Features
241
245
 
242
246
  - **Automatic prompt optimization** (`prompt_tune`) — correct a small sample in a browser UI, and the system generates per-category instructions that improve accuracy
@@ -97,7 +97,7 @@ cat.classify(
97
97
  ```
98
98
 
99
99
  ### `extract()`
100
- Discover categories from a corpus using LLM-driven exploration.
100
+ Discover categories from a corpus using LLM-driven exploration. Since v2.5.0, text consolidation runs the full explore → `collapse_themes()` pipeline (`engine="collapse"`, the default): the entire raw label inventory reaches the semantic merge — Jaro-Winkler dedup, embedding pre-merge, quality-controlled LLM passes, then a count-guided reduction to at most `max_categories`. Pass `engine="legacy"` to reproduce pre-2.5 runs (single merge call over a truncated inventory). `collapse_kwargs` forwards options to `collapse_themes()`; `max_workers` parallelizes both extraction and consolidation.
101
101
 
102
102
  ```python
103
103
  cat.extract(
@@ -156,7 +156,7 @@ cat.collapse_themes(
156
156
  | Parameter | Default | Description |
157
157
  | --- | --- | --- |
158
158
  | `input_data` | — | List of category labels, or a frequency `Series`/`dict` (`label -> count`). |
159
- | `api_key` | `None` | API key for the LLM provider (required). |
159
+ | `api_key` | `None` | API key for the LLM provider. Not required for subscription/CLI backends (`claude-code`, `claude-agent`, `codex-agent`) or `ollama`. |
160
160
  | `description` | `""` | The survey question or context, used in the merge prompt. |
161
161
  | `passes` | `1` | Number of merge iterations, or `"auto"` to iterate until the embedding-quality benchmark peaks. |
162
162
  | `max_passes` | `10` | Cap on iterations when `passes="auto"`. |
@@ -166,6 +166,8 @@ cat.collapse_themes(
166
166
  | `embedding_merge_threshold` | `0.92` | Cosine similarity at/above which labels are merged in the pre-LLM embedding step. `None`/`>=1.0` disables it. |
167
167
  | `shuffle` | `True` | Randomize order each pass so batch composition varies (improves convergence stability). |
168
168
  | `final_consolidation` | `0.82` | Cosine threshold for one greedy global embedding re-merge after all passes, collapsing cross-batch duplicates. Conservative by design (errs toward keeping categories). `False`/`None` skips it. |
169
+ | `top_n` | `None` | If set, a final global LLM call consolidates the surviving list into at most N categories, guided by each label's generation count (frequent themes favored; overlapping labels merged rather than dropped). Guaranteed `<= top_n` via truncation plus a deterministic top-N-by-count fallback. |
170
+ | `prune` | `False` | `True` = drop conceptual duplicates keeping one representative verbatim; never renames or merges merely-related labels. |
169
171
  | `user_model` | `"gpt-4o"` | Model for the merge phase. Use a capable model — small models can degenerate. |
170
172
  | `model_source` | `"auto"` | Provider for `user_model` (`"auto"`, `"openai"`, `"huggingface"`, …). |
171
173
  | `unique_model` | `None` | If set, run an initial extract-unique thinning phase on this (typically cheaper) model before the merge phase. `None` skips the phase (backward compatible). |
@@ -196,6 +198,8 @@ OpenAI, Anthropic, Google (Gemini), Mistral, Perplexity, xAI (Grok), HuggingFace
196
198
 
197
199
  All providers use the same `(model_name, provider, api_key)` tuple format. Provider is auto-detected from model name if omitted.
198
200
 
201
+ **Subscription backends (no API key).** Three `model_source` values authenticate through a chat subscription instead of a metered key — leave `api_key` unset: `"claude-agent"` (Claude subscription via the Agent SDK; `pip install "cat-stack[agent]"`), `"claude-code"` (the Claude Code CLI, if installed — no extra needed), and `"codex-agent"` (ChatGPT subscription; `pip install "cat-stack[codex-agent]"`). Classification, extraction, exploration, and summarization all route through them.
202
+
199
203
  ## Features
200
204
 
201
205
  - **Automatic prompt optimization** (`prompt_tune`) — correct a small sample in a browser UI, and the system generates per-category instructions that improve accuracy
@@ -42,8 +42,8 @@ embeddings = ["sentence-transformers>=2.2.0"]
42
42
  # `agent` keeps its historical meaning (the Claude backend) so every shipped
43
43
  # install hint stays true; `codex-agent` matches the provider string so the
44
44
  # error-message hint is copy-pasteable.
45
- agent = ["cat-claws[claude]>=0.3.0"]
46
- codex-agent = ["cat-claws[codex]>=0.3.0"]
45
+ agent = ["cat-claws[claude]>=0.3.1"]
46
+ codex-agent = ["cat-claws[codex]>=0.3.1"]
47
47
 
48
48
  [project.urls]
49
49
  Documentation = "https://github.com/chrissoria/cat-stack#readme"
@@ -1,7 +1,7 @@
1
1
  # SPDX-FileCopyrightText: 2025-present Christopher Soria <chrissoria@berkeley.edu>
2
2
  #
3
3
  # SPDX-License-Identifier: GPL-3.0-or-later
4
- __version__ = "2.4.0"
4
+ __version__ = "2.5.1"
5
5
  __author__ = "Chris Soria"
6
6
  __email__ = "chrissoria@berkeley.edu"
7
7
  __title__ = "cat-stack"
@@ -25,7 +25,7 @@ import numpy as np
25
25
  import pandas as pd
26
26
  from jellyfish import jaro_winkler_similarity
27
27
 
28
- from ._providers import UnifiedLLMClient, detect_provider
28
+ from ._providers import UnifiedLLMClient, detect_provider, _SUBSCRIPTION_PROVIDERS
29
29
  from ._utils import _clean_label
30
30
 
31
31
  __all__ = [
@@ -122,13 +122,18 @@ def _collapse_batch(client, batch, description, creativity, mode="unique"):
122
122
 
123
123
  mode="unique": extract unique categories only (remove restatements, keep
124
124
  distinct ones) — gentle, near-idempotent, guaranteed to only remove.
125
+ mode="prune": drop CONCEPTUAL duplicates (labels naming the same concept in
126
+ different words), keeping one representative per concept VERBATIM — never
127
+ renames or merges merely-related labels. Stronger than "unique" (catches
128
+ same-concept-different-words) but, unlike "merge", never collapses distinct
129
+ concepts. Output forced to a subset of the input.
125
130
  mode="merge": aggressively consolidate related labels into broader concepts
126
131
  while retaining meaningful distinctions — for a final compression step.
127
132
 
128
133
  Strict numbered-list prompt + strict parsing, so the reply is always a clean
129
134
  list and any stray prose is ignored. Guardrails: a failed call returns the
130
- batch unchanged (no data loss); in "unique" mode the output is forced to be a
131
- subset of the input (monotone, drift-free).
135
+ batch unchanged (no data loss); in "unique"/"prune" modes the output is forced
136
+ to be a subset of the input (monotone, drift-free).
132
137
  """
133
138
  items_blob = "; ".join(batch)
134
139
  context = f' about: "{description}"' if description else ""
@@ -150,6 +155,23 @@ def _collapse_batch(client, batch, description, creativity, mode="unique"):
150
155
  "2. Education\n"
151
156
  "3. Religion"
152
157
  )
158
+ elif mode == "prune":
159
+ prompt = (
160
+ f"You are given a list of category labels{context}. Some labels are DUPLICATES — they "
161
+ "name the SAME concept as another label using different words. Return the list with "
162
+ "duplicates removed. Rules: keep ONE label per distinct concept, copied EXACTLY as "
163
+ "written (verbatim); do NOT rename, rephrase, merge, or broaden any label; keep two "
164
+ "labels SEPARATE whenever they name genuinely different concepts, even if related; if "
165
+ "no labels duplicate each other, return them ALL unchanged. "
166
+ f"Labels are separated by semicolons within triple backticks: ```{items_blob}```.\n\n"
167
+ "Return ONLY a numbered list, using the labels exactly as they appear. Each line must "
168
+ "follow this exact format, with no other text before or after the list:\n"
169
+ "N. label\n\n"
170
+ "Example:\n"
171
+ "1. Employment\n"
172
+ "2. Education\n"
173
+ "3. Religion"
174
+ )
153
175
  else:
154
176
  prompt = (
155
177
  f"You are given a list of category labels{context}. "
@@ -186,8 +208,16 @@ def _collapse_batch(client, batch, description, creativity, mode="unique"):
186
208
  if label:
187
209
  out.append(label)
188
210
 
189
- if mode == "unique":
190
- # Contraction guarantee: extract-unique must only REMOVE, never add or
211
+ if not out:
212
+ # Unparseable reply (no numbered list found): keep the batch unchanged
213
+ # rather than silently dropping every label in it. Mirrors the API-error
214
+ # fallback above; matters most in "merge" mode, which has no
215
+ # subset-of-input guarantee to fall back on.
216
+ sys.stderr.write("[collapse_themes] unparseable reply — keeping batch unchanged\n")
217
+ return [str(x).strip().lower() for x in batch]
218
+
219
+ if mode in ("unique", "prune"):
220
+ # Contraction guarantee: extract-unique/prune must only REMOVE, never add or
191
221
  # mutate. Keep only outputs that map back to an input label (by normalized
192
222
  # key), as the original input string. Makes every pass monotone and
193
223
  # drift-free, immune to intermittent model rephrasing/splitting.
@@ -205,6 +235,85 @@ def _collapse_batch(client, batch, description, creativity, mode="unique"):
205
235
  return out
206
236
 
207
237
 
238
+ def _select_top_n(client, counts, n, description, creativity, dedupe_threshold):
239
+ """The "nuclear" final step: one global LLM call that consolidates the whole
240
+ surviving list into at most `n` categories, guided by generation counts
241
+ (higher = the theme recurred more often during extraction). Overlapping
242
+ labels are merged rather than dropped; only rare, unabsorbable themes fall
243
+ away. The result is guaranteed <= n: over-length replies are truncated, and
244
+ a failed or unparseable call falls back to the top-n labels by count
245
+ (deterministic)."""
246
+ ordered = sorted(counts, key=counts.get, reverse=True)
247
+ fallback = [str(x).strip().lower() for x in ordered[:n]]
248
+ blob = "; ".join(f"{lbl} ({counts[lbl]})" for lbl in ordered)
249
+ context = f' about: "{description}"' if description else ""
250
+ prompt = (
251
+ f"You are given category labels{context}. Each label is followed by a count in "
252
+ "parentheses: how many times it was independently generated during extraction — "
253
+ "higher counts mean the theme is more common in the data. Consolidate this list "
254
+ f"into EXACTLY {n} final categories that best summarize the data. Favor "
255
+ "high-count themes; merge overlapping or related labels into one clear "
256
+ "representative label rather than dropping them; drop only rare themes that "
257
+ f"cannot reasonably fold into any of the {n}. Labels are separated by semicolons "
258
+ f"within triple backticks: ```{blob}```\n\n"
259
+ f"Return ONLY a numbered list of exactly {n} category labels, without counts. "
260
+ "Each line must follow this exact format, with no other text before or after "
261
+ "the list:\n"
262
+ "N. label\n\n"
263
+ "Example:\n"
264
+ "1. Employment\n"
265
+ "2. Education\n"
266
+ "3. Religion"
267
+ )
268
+ reply, error = client.complete(
269
+ messages=[{"role": "user", "content": prompt}],
270
+ creativity=creativity,
271
+ force_json=False,
272
+ )
273
+ if error:
274
+ sys.stderr.write(f"[collapse_themes] top_n call failed: {error} — "
275
+ "falling back to top-n by count\n")
276
+ return fallback
277
+ out = []
278
+ for line in (reply or "").splitlines():
279
+ m = _LINE_PAT.match(line.strip())
280
+ if m:
281
+ label = _clean_label(m.group(1)).strip(" ;.,")
282
+ if label:
283
+ out.append(label)
284
+ out = _jw_dedupe(out, dedupe_threshold)
285
+ if not out:
286
+ sys.stderr.write("[collapse_themes] top_n reply unparseable — "
287
+ "falling back to top-n by count\n")
288
+ return fallback
289
+ return out[:n]
290
+
291
+
292
+ def _count_guidance(current, input_data):
293
+ """{label: count} evidence for the top_n prompt: each surviving label gets its
294
+ aggregate generation count from the ORIGINAL input (matched by normalized key)
295
+ when available, else its count in the current list."""
296
+ cur = _to_counts(current)
297
+ orig = {}
298
+ for k, v in _to_counts(input_data).items():
299
+ key = _norm_key(k)
300
+ orig[key] = orig.get(key, 0) + int(v)
301
+ return {lbl: max(int(c), orig.get(_norm_key(lbl), 0)) for lbl, c in cur.items()}
302
+
303
+
304
+ def _as_label_list(x):
305
+ """Enforce the documented list[str] return contract at exit points.
306
+
307
+ Dict / Series / DataFrame inputs can reach a return untouched (e.g.
308
+ passes=0 with a top_n no-op), which would leak the input container —
309
+ and, via the filename write, save counts instead of labels."""
310
+ if isinstance(x, list):
311
+ return [str(v) for v in x]
312
+ if isinstance(x, dict):
313
+ return [str(k) for k in x]
314
+ return [str(k) for k in _to_counts(x)]
315
+
316
+
208
317
  def _to_counts(input_data):
209
318
  """Coerce the accepted input forms into a {category: count} dict."""
210
319
  if isinstance(input_data, pd.DataFrame):
@@ -288,6 +397,9 @@ def collapse_themes(
288
397
  embedding_merge_threshold=0.92,
289
398
  shuffle=True,
290
399
  final_consolidation=0.82,
400
+ top_n=None,
401
+ prune=False,
402
+ prune_threshold=50,
291
403
  user_model="gpt-4o",
292
404
  model_source="auto",
293
405
  unique_model=None,
@@ -326,13 +438,17 @@ def collapse_themes(
326
438
  input_data: Themes to collapse. list[str] (duplicates allowed), pandas
327
439
  Series, dict {category: count}, or DataFrame with "category"
328
440
  [and optional "count"] columns.
329
- api_key (str): API key for the model provider.
441
+ api_key (str): API key for the model provider. Not required for
442
+ subscription/CLI backends (`claude-code`, `claude-agent`,
443
+ `codex-agent`) or `ollama`.
330
444
  description (str): Data/question context, injected into the prompt — e.g.
331
445
  the survey question the categories came from. Helps the model judge
332
446
  which distinctions matter.
333
447
  passes (int | str): Number of collapse iterations, or "auto" to iterate
334
448
  until the deterministic quality benchmark peaks (the recommended mode
335
- for a final taxonomy — pair with aggressive=True). Default 1.
449
+ for a final taxonomy — pair with aggressive=True). Either way,
450
+ iteration stops early once two consecutive passes leave the list
451
+ unchanged. Default 1.
336
452
  max_passes (int): Cap on iterations when passes="auto". Default 10.
337
453
  batch_size (int): Themes per LLM chunk (ceil(n / batch_size) calls per
338
454
  pass). Default 40.
@@ -345,12 +461,29 @@ def collapse_themes(
345
461
  0.92. None or >=1.0 skips embeddings.
346
462
  shuffle (bool): Randomize order each pass so batch composition varies.
347
463
  Default True (improves convergence stability).
464
+ prune (bool): If True, use the prune strategy instead of the merge/unique
465
+ passes: drop only CONCEPTUAL duplicates (one representative per concept,
466
+ kept verbatim), never renaming or merging merely-related labels, so a
467
+ short clean list is left intact while a long list is deduplicated. Lists
468
+ at/below prune_threshold go straight to a single global prune; longer
469
+ lists are reduced by batched prune first, then a final global prune.
470
+ prune_threshold (int): Max list length that goes directly to a single global
471
+ prune call (no batching). Default 50.
348
472
  final_consolidation (float): Cosine threshold for one greedy embedding
349
473
  re-merge over the whole result after all passes, collapsing cross-batch
350
474
  lexical-sibling duplicates that batched passes (and the auto loop) cannot
351
475
  reach. Default 0.82 — deterministic and tuned to land just above the true
352
476
  concept count (errs toward keeping categories; over-segmentation is
353
477
  preferred over over-consolidation). False/None skips.
478
+ top_n (int): Optional "nuclear" final step. After all passes and the final
479
+ consolidation, ONE global LLM call consolidates the surviving list into
480
+ at most `top_n` categories, guided by each label's generation count
481
+ (frequent themes favored; overlapping labels merged into one
482
+ representative rather than dropped). None (default) changes nothing
483
+ about existing behavior. The result is guaranteed to have <= top_n
484
+ labels: a failed or unparseable call falls back deterministically to
485
+ the top_n labels by count. Use when a compact fixed-size taxonomy is
486
+ required and coverage of rare themes is knowingly traded away.
354
487
  user_model (str): Model name for the merge phase. Default "gpt-4o". Use a
355
488
  capable model — small models can degenerate into repetition.
356
489
  model_source (str): Provider — "auto", "openai", "huggingface", etc.
@@ -378,7 +511,7 @@ def collapse_themes(
378
511
  list[str]: The collapsed category list after `passes` iterations.
379
512
 
380
513
  Examples:
381
- >>> import cat_stack as cat
514
+ >>> import catstack as cat
382
515
  >>> themes = cat.explore(df['responses'], description="Why did you move?",
383
516
  ... api_key=key)
384
517
  >>> # Recommended: aggressive merge, auto-stop at the quality peak
@@ -387,9 +520,6 @@ def collapse_themes(
387
520
  ... aggressive=True, passes="auto", max_workers=8,
388
521
  ... )
389
522
  """
390
- if not api_key:
391
- raise ValueError("collapse_themes() needs an api_key for the LLM call.")
392
-
393
523
  mode = "merge" if aggressive else "unique"
394
524
 
395
525
  # The main (merge) phase runs on merge_model if given, else user_model. A separate
@@ -400,6 +530,20 @@ def collapse_themes(
400
530
  merge_name = merge_model or user_model
401
531
  merge_src = merge_model_source if merge_model else model_source
402
532
  merge_provider = detect_provider(merge_name, merge_src)
533
+
534
+ # A key is only needed for providers that bill one; the subscription/CLI
535
+ # backends and ollama run keyless.
536
+ _keyless = set(_SUBSCRIPTION_PROVIDERS) | {"ollama"}
537
+ _providers_used = {merge_provider}
538
+ if unique_model:
539
+ _providers_used.add(detect_provider(unique_model, unique_model_source))
540
+ if not api_key and not _providers_used <= _keyless:
541
+ raise ValueError(
542
+ "collapse_themes() needs an api_key for the LLM call. "
543
+ "(Not required for the claude-code/claude-agent/codex-agent "
544
+ "backends or ollama.)"
545
+ )
546
+
403
547
  client = UnifiedLLMClient(provider=merge_provider, api_key=api_key, model=merge_name)
404
548
 
405
549
  def _run(cl, items, md, p):
@@ -419,6 +563,64 @@ def collapse_themes(
419
563
  def _pass(items, p):
420
564
  return _run(client, items, mode, p)
421
565
 
566
+ # ── Prune strategy ───────────────────────────────────────────────────────
567
+ # Drop CONCEPTUAL duplicates only (keep one representative per concept,
568
+ # verbatim) — never rename or merge merely-related labels, so it cannot
569
+ # over-consolidate a short, already-clean list. Length-routed: a list at or
570
+ # below prune_threshold goes straight to a single global "master list" prune;
571
+ # a longer list is first reduced by the same prune primitive in shuffled
572
+ # batches (until it stops shrinking or fits the threshold), then finished with
573
+ # the global master prune.
574
+ if prune:
575
+ items = list(_to_counts(input_data).keys())
576
+ items = _jw_dedupe(items, dedupe_threshold) # gentle, no rename
577
+ items = _embedding_merge(items, embedding_merge_threshold)
578
+ rng = random.Random(random_state)
579
+ guard = 0
580
+ while len(items) > prune_threshold and guard < max_passes:
581
+ guard += 1
582
+ if shuffle:
583
+ rng.shuffle(items)
584
+ batches = [items[i:i + prune_threshold]
585
+ for i in range(0, len(items), prune_threshold)]
586
+ if max_workers and max_workers > 1 and len(batches) > 1:
587
+ from concurrent.futures import ThreadPoolExecutor, as_completed
588
+
589
+ results = [None] * len(batches)
590
+ with ThreadPoolExecutor(max_workers=max_workers) as ex:
591
+ futures = {
592
+ ex.submit(_collapse_batch, client, b, description,
593
+ creativity, "prune"): i
594
+ for i, b in enumerate(batches)
595
+ }
596
+ for fut in as_completed(futures):
597
+ results[futures[fut]] = fut.result()
598
+ out = [label for r in results for label in (r or [])]
599
+ else:
600
+ out = []
601
+ for b in batches:
602
+ out += _collapse_batch(client, b, description, creativity,
603
+ mode="prune")
604
+ out = _jw_dedupe(out, dedupe_threshold)
605
+ if progress_callback:
606
+ progress_callback(guard, max_passes, "collapse_themes:prune")
607
+ if len(out) >= len(items): # no further reduction this round
608
+ items = out
609
+ break
610
+ items = out
611
+ if len(items) > 1: # final global master-list prune
612
+ items = _jw_dedupe(
613
+ _collapse_batch(client, items, description, creativity, mode="prune"),
614
+ dedupe_threshold)
615
+ if top_n and len(items) > int(top_n):
616
+ items = _select_top_n(client, _count_guidance(items, input_data),
617
+ int(top_n), description, creativity, dedupe_threshold)
618
+ items = _as_label_list(items)
619
+ if filename:
620
+ pd.DataFrame({"category": items}).to_csv(filename, index=False)
621
+ print(f"Collapsed categories saved to {filename}")
622
+ return items
623
+
422
624
  current = input_data
423
625
 
424
626
  # Phase 1 (optional): cheap unique-keeping thin. When unique_model is set, run
@@ -442,6 +644,7 @@ def collapse_themes(
442
644
  show_progress_bar=False,
443
645
  )
444
646
  best, best_q = None, -1.0
647
+ prev_key, stable = None, 0
445
648
  for p in range(max_passes):
446
649
  current = _pass(current, p)
447
650
  q = _quality(current, raw_embs)
@@ -450,12 +653,26 @@ def collapse_themes(
450
653
  if q < best_q:
451
654
  break # quality dropped -> the previous pass was the peak
452
655
  best, best_q = current, q
656
+ key = tuple(sorted(current))
657
+ stable = stable + 1 if key == prev_key else 0
658
+ prev_key = key
659
+ if stable >= 2:
660
+ break # converged: two consecutive passes changed nothing
453
661
  current = best if best is not None else current
454
662
  else:
663
+ prev_key, stable = None, 0
455
664
  for p in range(int(passes)):
456
665
  current = _pass(current, p)
457
666
  if progress_callback:
458
667
  progress_callback(p + 1, int(passes), "collapse_themes")
668
+ # Early stop at a fixed point: shuffling gives labels one fresh batch
669
+ # composition to merge under; if two consecutive passes both change
670
+ # nothing, further passes are near-certain no-ops — stop burning calls.
671
+ key = tuple(sorted(current))
672
+ stable = stable + 1 if key == prev_key else 0
673
+ prev_key = key
674
+ if stable >= 2:
675
+ break
459
676
 
460
677
  # Final global consolidation. Batched passes (and the auto loop) can only merge
461
678
  # labels that share a batch, so cross-batch lexical siblings — e.g. "tension" vs
@@ -470,8 +687,21 @@ def collapse_themes(
470
687
  # categories — over-segmentation is the preferred failure mode, not
471
688
  # over-consolidation. Set final_consolidation=False to skip.
472
689
  if final_consolidation and len(current) > 1:
690
+ # Order by original generation frequency first, so the greedy merge keeps
691
+ # the most frequently generated variant as each cluster's representative
692
+ # instead of whichever label the last (shuffled) pass happened to emit
693
+ # first. Labels the merge phase renamed have no original count and sort
694
+ # last, preserving their relative order (sort is stable).
695
+ orig_counts = {_norm_key(k): v for k, v in _to_counts(input_data).items()}
696
+ current = sorted(current, key=lambda c: orig_counts.get(_norm_key(c), 0),
697
+ reverse=True)
473
698
  current = _embedding_merge(current, final_consolidation)
474
699
 
700
+ if top_n and len(_to_counts(current)) > int(top_n):
701
+ current = _select_top_n(client, _count_guidance(current, input_data),
702
+ int(top_n), description, creativity, dedupe_threshold)
703
+
704
+ current = _as_label_list(current)
475
705
  if filename:
476
706
  pd.DataFrame({"category": current}).to_csv(filename, index=False)
477
707
  print(f"Collapsed categories saved to {filename}")
@@ -16,7 +16,7 @@ from .text_functions import explore_common_categories
16
16
 
17
17
  def explore(
18
18
  input_data,
19
- api_key,
19
+ api_key=None,
20
20
  description="",
21
21
  max_categories=12,
22
22
  categories_per_chunk=10,
@@ -46,7 +46,9 @@ def explore(
46
46
 
47
47
  Args:
48
48
  input_data: List of text responses or pandas Series.
49
- api_key (str): API key for the model provider.
49
+ api_key (str, optional): API key for the model provider. Not required
50
+ for subscription/CLI backends (claude-code, claude-agent,
51
+ codex-agent) or ollama.
50
52
  description (str): Description of the data context. Content-neutral —
51
53
  for survey responses this is the question that was asked; for
52
54
  documents or posts this describes what the content is about.
@@ -77,7 +79,7 @@ def explore(
77
79
  every iteration. Length ≈ iterations × divisions × categories_per_chunk.
78
80
 
79
81
  Examples:
80
- >>> import cat_stack as cat
82
+ >>> import catstack as cat
81
83
  >>>
82
84
  >>> raw_categories = cat.explore(
83
85
  ... input_data=df['responses'],
@@ -36,10 +36,12 @@ from .pdf_functions import (
36
36
  explore_pdf_categories,
37
37
  )
38
38
 
39
+ from .collapse_themes import collapse_themes
40
+
39
41
 
40
42
  def extract(
41
43
  input_data,
42
- api_key,
44
+ api_key=None,
43
45
  input_type="auto",
44
46
  description="",
45
47
  survey_question=None,
@@ -61,6 +63,9 @@ def extract(
61
63
  auto_download: bool = False,
62
64
  input_mode=None,
63
65
  domain: str = "neutral",
66
+ engine: str = "collapse",
67
+ max_workers: int = 1,
68
+ collapse_kwargs: dict = None,
64
69
  ):
65
70
  """
66
71
  Unified category extraction function for text, image, and PDF inputs.
@@ -74,7 +79,9 @@ def extract(
74
79
  - For text: list of text responses or pandas Series
75
80
  - For image: directory path, single file, or list of image paths
76
81
  - For pdf: directory path, single file, or list of PDF paths
77
- api_key (str): API key for the model provider.
82
+ api_key (str, optional): API key for the model provider. Not required
83
+ for subscription/CLI backends (claude-code, claude-agent,
84
+ codex-agent) or ollama.
78
85
  input_type (str): Type of input data. Options:
79
86
  - "auto" (default): Auto-detect from file extensions
80
87
  - "text": Text responses
@@ -111,15 +118,31 @@ def extract(
111
118
  limits. Default 0.0 (no delay).
112
119
  auto_download (bool): If True, automatically download missing Ollama
113
120
  models without prompting. Default False.
121
+ engine (str): Consolidation engine for text input. "collapse" (default):
122
+ run raw extraction (as explore() does) and consolidate the FULL
123
+ inventory with collapse_themes() — semantic pre-clean, quality-
124
+ controlled passes, then a count-guided reduction to at most
125
+ `max_categories`. "legacy": the pre-2.5 single merge call, which
126
+ truncates the inventory to the top max_categories*3 labels by
127
+ exact-string count before merging; kept for reproducing older runs.
128
+ max_workers (int): Parallel API calls for both the extraction chunks and
129
+ the consolidation batches (text engine="collapse" only; extraction
130
+ also honors it under "legacy"). Default 1.
131
+ collapse_kwargs (dict): Optional overrides forwarded to collapse_themes()
132
+ when engine="collapse" — e.g. {"prune": True} or
133
+ {"passes": 2, "aggressive": False}. Defaults applied first:
134
+ passes="auto", aggressive=True. `top_n` cannot be overridden here;
135
+ it is always max_categories.
114
136
 
115
137
  Returns:
116
138
  dict with keys:
117
139
  - counts_df: DataFrame of categories with counts
118
140
  - top_categories: List of top category names
119
- - raw_top_text: Raw model output from final merge step
141
+ - raw_top_text: Raw model output from final merge step ("" when
142
+ engine="collapse", which has no single merge reply)
120
143
 
121
144
  Examples:
122
- >>> import cat_stack as cat
145
+ >>> import catstack as cat
123
146
  >>>
124
147
  >>> # Extract categories from text responses
125
148
  >>> results = cat.extract(
@@ -171,7 +194,41 @@ def extract(
171
194
  resolved_description = description or ""
172
195
 
173
196
  if input_type == "text":
174
- return explore_common_categories(
197
+ if engine not in ("collapse", "legacy"):
198
+ raise ValueError(f"engine must be 'collapse' or 'legacy', got '{engine}'")
199
+
200
+ if engine == "legacy":
201
+ return explore_common_categories(
202
+ input_data=input_data,
203
+ api_key=api_key,
204
+ survey_question=resolved_description,
205
+ max_categories=max_categories,
206
+ categories_per_chunk=categories_per_chunk,
207
+ divisions=divisions,
208
+ user_model=user_model,
209
+ creativity=creativity,
210
+ specificity=specificity,
211
+ research_question=research_question,
212
+ filename=filename,
213
+ model_source=model_source,
214
+ iterations=iterations,
215
+ random_state=random_state,
216
+ focus=focus,
217
+ progress_callback=progress_callback,
218
+ chunk_delay=chunk_delay,
219
+ auto_download=auto_download,
220
+ max_workers=max_workers,
221
+ domain=domain,
222
+ )
223
+
224
+ # engine="collapse": raw extraction (what explore() does), then consolidate
225
+ # the FULL inventory with collapse_themes(). Unlike the legacy merge, no
226
+ # label is truncated away before consolidation, pre-cleaning is semantic
227
+ # (Jaro-Winkler + embeddings) rather than exact-string, and the final
228
+ # count-guided top_n step guarantees at most max_categories categories.
229
+ import pandas as pd
230
+
231
+ raw_items = explore_common_categories(
175
232
  input_data=input_data,
176
233
  api_key=api_key,
177
234
  survey_question=resolved_description,
@@ -182,17 +239,63 @@ def extract(
182
239
  creativity=creativity,
183
240
  specificity=specificity,
184
241
  research_question=research_question,
185
- filename=filename,
242
+ filename=None,
186
243
  model_source=model_source,
187
244
  iterations=iterations,
188
245
  random_state=random_state,
189
246
  focus=focus,
190
247
  progress_callback=progress_callback,
248
+ return_raw=True,
191
249
  chunk_delay=chunk_delay,
192
250
  auto_download=auto_download,
251
+ max_workers=max_workers,
193
252
  domain=domain,
194
253
  )
195
254
 
255
+ # Frequency inventory in the same shape the legacy engine returned.
256
+ def _normalize(cat):
257
+ return "/".join(sorted(t.strip().lower() for t in str(cat).split("/")))
258
+
259
+ flat = [str(x).strip() for x in raw_items if str(x).strip()]
260
+ if not flat:
261
+ raise ValueError("No categories were extracted from the model responses.")
262
+ inv = pd.DataFrame(flat, columns=["Category"])
263
+ inv["normalized"] = inv["Category"].map(_normalize)
264
+ counts_df = (
265
+ inv.groupby("normalized")
266
+ .agg(Category=("Category", lambda x: x.value_counts().index[0]),
267
+ counts=("Category", "size"))
268
+ .sort_values("counts", ascending=False)
269
+ .reset_index(drop=True)
270
+ )
271
+
272
+ ck = dict(passes="auto", aggressive=True)
273
+ ck.update(collapse_kwargs or {})
274
+ ck["top_n"] = int(max_categories) # the required N — not overridable
275
+ top = collapse_themes(
276
+ raw_items,
277
+ api_key=api_key,
278
+ description=resolved_description,
279
+ user_model=user_model,
280
+ model_source=model_source,
281
+ creativity=0 if creativity is None else creativity,
282
+ max_workers=max_workers,
283
+ random_state=random_state,
284
+ progress_callback=progress_callback,
285
+ **ck,
286
+ )
287
+
288
+ if filename:
289
+ pd.DataFrame({"rank": range(1, len(top) + 1), "category": top}).to_csv(
290
+ filename, index=False)
291
+ print(f"Top {len(top)} categories saved to {filename}")
292
+
293
+ return {
294
+ "counts_df": counts_df,
295
+ "top_categories": top,
296
+ "raw_top_text": "",
297
+ }
298
+
196
299
  elif input_type == "image":
197
300
  return explore_image_categories(
198
301
  image_input=input_data,
File without changes
File without changes