cat-stack 2.0.0b1__tar.gz → 2.0.0b5__tar.gz

{cat_stack-2.0.0b1 → cat_stack-2.0.0b5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cat-stack
-Version: 2.0.0b1
+Version: 2.0.0b5
 Summary: Domain-agnostic text, image, PDF, and DOCX classification engine powered by LLMs
 Project-URL: Documentation, https://github.com/chrissoria/cat-stack#readme
 Project-URL: Issues, https://github.com/chrissoria/cat-stack/issues
@@ -154,6 +154,67 @@ cat.explore(
 )
 ```
 
+### `collapse_themes()`
+Consolidate a long, redundant list of extracted category labels (e.g. the output of `explore()`) into a smaller, deduplicated taxonomy. Runs the semantic merge iteratively, then applies a single deterministic embedding re-merge over the whole result to collapse cross-batch lexical siblings (e.g. "tension" / "estrangement") that batched passes leave separate. Tuned to err toward over-segmentation (keeping categories) rather than over-merging.
+
+```python
+# Basic: aggressive merge, auto-stop at the quality peak
+cat.collapse_themes(
+    input_data=raw_labels,            # list[str] or a frequency Series/dict
+    api_key=key,
+    description="Why did you move?",   # the survey question / context
+    aggressive=True,
+    passes="auto",
+    user_model="gpt-4o",
+)
+```
+
+```python
+# Per-step model assignment: a cheap model thins restatements,
+# a stronger model does the conceptual merge (providers can differ)
+cat.collapse_themes(
+    input_data=raw_labels,
+    api_key=key,
+    description="Why did you move?",
+    aggressive=True,
+    passes="auto",
+    unique_model="Qwen/Qwen2.5-72B-Instruct:together",
+    unique_model_source="huggingface",
+    unique_passes=1,
+    merge_model="Qwen/Qwen3.6-35B-A3B:together",
+    merge_model_source="huggingface",
+    max_workers=8,
+)
+```
+
+**Parameters**
+
+| Parameter | Default | Description |
+| --- | --- | --- |
+| `input_data` | — | List of category labels, or a frequency `Series`/`dict` (`label -> count`). |
+| `api_key` | `None` | API key for the LLM provider (required). |
+| `description` | `""` | The survey question or context, used in the merge prompt. |
+| `passes` | `1` | Number of merge iterations, or `"auto"` to iterate until the embedding-quality benchmark peaks. |
+| `max_passes` | `10` | Cap on iterations when `passes="auto"`. |
+| `batch_size` | `40` | Labels per LLM chunk (`ceil(n / batch_size)` calls per pass). |
+| `aggressive` | `False` | `True` = conceptual-merge prompt (compress related labels); `False` = extract-unique (faithful thinning, removes restatements only). |
+| `dedupe_threshold` | `0.95` | Jaro-Winkler similarity at/above which normalized labels are deduped (`1.0` = exact only). |
+| `embedding_merge_threshold` | `0.92` | Cosine similarity at/above which labels are merged in the pre-LLM embedding step. `None`/`>=1.0` disables it. |
+| `shuffle` | `True` | Randomize order each pass so batch composition varies (improves convergence stability). |
+| `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. |
+| `user_model` | `"gpt-4o"` | Model for the merge phase. Use a capable model — small models can degenerate. |
+| `model_source` | `"auto"` | Provider for `user_model` (`"auto"`, `"openai"`, `"huggingface"`, …). |
+| `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). |
+| `unique_model_source` | `"auto"` | Provider for `unique_model` — can differ from the merge phase. |
+| `unique_passes` | `1` | Number of thinning passes when `unique_model` is set. |
+| `merge_model` | `None` | Model for the merge phase; falls back to `user_model` when `None`. |
+| `merge_model_source` | `"auto"` | Provider for `merge_model`. |
+| `creativity` | `0` | Temperature (`0` = deterministic). |
+| `max_workers` | `1` | Batches processed concurrently per pass. |
+| `random_state` | `None` | Seed for shuffling (per-pass seed = `random_state + pass`). |
+| `filename` | `None` | Optional CSV path to save the final list. |
+| `progress_callback` | `None` | Optional `callback(pass, passes, label)` for progress reporting. |
+
 ### `summarize()`
 Summarize text or PDF documents, with optional multi-model ensemble.
 

{cat_stack-2.0.0b1 → cat_stack-2.0.0b5}/README.md

@@ -118,6 +118,67 @@ cat.explore(
 )
 ```
 
+### `collapse_themes()`
+Consolidate a long, redundant list of extracted category labels (e.g. the output of `explore()`) into a smaller, deduplicated taxonomy. Runs the semantic merge iteratively, then applies a single deterministic embedding re-merge over the whole result to collapse cross-batch lexical siblings (e.g. "tension" / "estrangement") that batched passes leave separate. Tuned to err toward over-segmentation (keeping categories) rather than over-merging.
+
+```python
+# Basic: aggressive merge, auto-stop at the quality peak
+cat.collapse_themes(
+    input_data=raw_labels,            # list[str] or a frequency Series/dict
+    api_key=key,
+    description="Why did you move?",   # the survey question / context
+    aggressive=True,
+    passes="auto",
+    user_model="gpt-4o",
+)
+```
+
+```python
+# Per-step model assignment: a cheap model thins restatements,
+# a stronger model does the conceptual merge (providers can differ)
+cat.collapse_themes(
+    input_data=raw_labels,
+    api_key=key,
+    description="Why did you move?",
+    aggressive=True,
+    passes="auto",
+    unique_model="Qwen/Qwen2.5-72B-Instruct:together",
+    unique_model_source="huggingface",
+    unique_passes=1,
+    merge_model="Qwen/Qwen3.6-35B-A3B:together",
+    merge_model_source="huggingface",
+    max_workers=8,
+)
+```
+
+**Parameters**
+
+| Parameter | Default | Description |
+| --- | --- | --- |
+| `input_data` | — | List of category labels, or a frequency `Series`/`dict` (`label -> count`). |
+| `api_key` | `None` | API key for the LLM provider (required). |
+| `description` | `""` | The survey question or context, used in the merge prompt. |
+| `passes` | `1` | Number of merge iterations, or `"auto"` to iterate until the embedding-quality benchmark peaks. |
+| `max_passes` | `10` | Cap on iterations when `passes="auto"`. |
+| `batch_size` | `40` | Labels per LLM chunk (`ceil(n / batch_size)` calls per pass). |
+| `aggressive` | `False` | `True` = conceptual-merge prompt (compress related labels); `False` = extract-unique (faithful thinning, removes restatements only). |
+| `dedupe_threshold` | `0.95` | Jaro-Winkler similarity at/above which normalized labels are deduped (`1.0` = exact only). |
+| `embedding_merge_threshold` | `0.92` | Cosine similarity at/above which labels are merged in the pre-LLM embedding step. `None`/`>=1.0` disables it. |
+| `shuffle` | `True` | Randomize order each pass so batch composition varies (improves convergence stability). |
+| `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. |
+| `user_model` | `"gpt-4o"` | Model for the merge phase. Use a capable model — small models can degenerate. |
+| `model_source` | `"auto"` | Provider for `user_model` (`"auto"`, `"openai"`, `"huggingface"`, …). |
+| `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). |
+| `unique_model_source` | `"auto"` | Provider for `unique_model` — can differ from the merge phase. |
+| `unique_passes` | `1` | Number of thinning passes when `unique_model` is set. |
+| `merge_model` | `None` | Model for the merge phase; falls back to `user_model` when `None`. |
+| `merge_model_source` | `"auto"` | Provider for `merge_model`. |
+| `creativity` | `0` | Temperature (`0` = deterministic). |
+| `max_workers` | `1` | Batches processed concurrently per pass. |
+| `random_state` | `None` | Seed for shuffling (per-pass seed = `random_state + pass`). |
+| `filename` | `None` | Optional CSV path to save the final list. |
+| `progress_callback` | `None` | Optional `callback(pass, passes, label)` for progress reporting. |
+
 ### `summarize()`
 Summarize text or PDF documents, with optional multi-model ensemble.
 

{cat_stack-2.0.0b1 → cat_stack-2.0.0b5}/src/catstack/__about__.py

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: 2025-present Christopher Soria <chrissoria@berkeley.edu>
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
-__version__ = "2.0.0b1"
+__version__ = "2.0.0b5"
 __author__ = "Chris Soria"
 __email__ = "chrissoria@berkeley.edu"
 __title__ = "cat-stack"

{cat_stack-2.0.0b1 → cat_stack-2.0.0b5}/src/catstack/_providers.py

@@ -333,6 +333,14 @@ def _detect_huggingface_endpoint(api_key: str, model: str, skip: set = None) ->
     headers = {
         "Content-Type": "application/json",
         "Authorization": f"Bearer {api_key}",
+        # Match the main request path: featherless's WAF 403s the default
+        # python-requests agent, which would make this probe wrongly skip a
+        # working endpoint.
+        "User-Agent": (
+            "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
+            "AppleWebKit/537.36 (KHTML, like Gecko) "
+            "Chrome/120.0.0.0 Safari/537.36"
+        ),
     }
     payload = {
         "model": clean_model,
@@ -526,7 +534,19 @@ class UnifiedLLMClient:
 
     def _get_headers(self) -> dict:
         """Build request headers for the provider."""
-        headers = {"Content-Type": "application/json"}
+        # Send a browser-like User-Agent. Some providers fronted by a WAF
+        # (notably the HuggingFace router's featherless-ai backend) intermittently
+        # 403 the default `python-requests/x.y` agent via a Cloudflare bot rule,
+        # which surfaces as spurious classification failures. A standard UA is
+        # accepted everywhere and costs nothing on providers that don't care.
+        headers = {
+            "Content-Type": "application/json",
+            "User-Agent": (
+                "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
+                "AppleWebKit/537.36 (KHTML, like Gecko) "
+                "Chrome/120.0.0.0 Safari/537.36"
+            ),
+        }
         auth_header = self.config["auth_header"]
         auth_prefix = self.config["auth_prefix"]
 

{cat_stack-2.0.0b1 → cat_stack-2.0.0b5}/src/catstack/collapse_themes.py

@@ -70,17 +70,21 @@ def _jw_dedupe(items, threshold):
     return out
 
 
-def _embedding_merge(items, threshold):
-    """Greedy embedding clustering: drop labels whose cosine similarity to an
-    already-kept label is >= threshold. Keeps the first-seen representative.
-    Uses cat-stack's canonical BAAI/bge-small model (cached)."""
+def _get_emb_model():
+    """Load (once) and return cat-stack's canonical BAAI/bge-small embedder."""
     global _EMB_MODEL
-    if not threshold or threshold >= 1.0 or len(items) < 2:
-        return items
     if _EMB_MODEL is None:
         from ._embeddings import load_embedding_model
         _EMB_MODEL = load_embedding_model()
-    embs = _EMB_MODEL.encode(items, normalize_embeddings=True, show_progress_bar=False)
+    return _EMB_MODEL
+
+
+def _embedding_merge(items, threshold):
+    """Greedy embedding clustering: drop labels whose cosine similarity to an
+    already-kept label is >= threshold. Keeps the first-seen representative."""
+    if not threshold or threshold >= 1.0 or len(items) < 2:
+        return items
+    embs = _get_emb_model().encode(items, normalize_embeddings=True, show_progress_bar=False)
     reps, rep_embs = [], []
     for it, e in zip(items, embs):
         if rep_embs and float(np.max(np.asarray(rep_embs) @ e)) >= threshold:
@@ -90,6 +94,29 @@ def _embedding_merge(items, threshold):
     return reps
 
 
+def _quality(output, raw_embs, tau_cov=0.70, tau_red=0.85, beta=2.0):
+    """Deterministic quality of a candidate taxonomy vs the raw input themes:
+    coverage-weighted F-beta of recall=coverage_hard (share of raw within tau_cov
+    of some output) and precision=(1 - redundancy_rate) (share of outputs with a
+    near-twin >= tau_red). Embedding-only — the convergence signal for passes='auto'.
+    """
+    if not output:
+        return 0.0
+    O = _get_emb_model().encode(list(output), normalize_embeddings=True, show_progress_bar=False)
+    coverage = float(((raw_embs @ O.T).max(axis=1) >= tau_cov).mean())
+    if len(output) > 1:
+        OO = O @ O.T
+        np.fill_diagonal(OO, -1.0)
+        redundancy = float((OO.max(axis=1) >= tau_red).mean())
+    else:
+        redundancy = 0.0
+    precision = 1.0 - redundancy
+    if coverage <= 0 or precision <= 0:
+        return 0.0
+    b2 = beta * beta
+    return (1 + b2) * precision * coverage / (b2 * precision + coverage)
+
+
 def _collapse_batch(client, batch, description, creativity, mode="unique"):
     """One LLM call on a single batch -> list[str].
 
@@ -254,13 +281,20 @@ def collapse_themes(
     api_key=None,
     description="",
     passes=1,
+    max_passes=10,
     batch_size=40,
     aggressive=False,
     dedupe_threshold=0.95,
     embedding_merge_threshold=0.92,
     shuffle=True,
+    final_consolidation=0.82,
     user_model="gpt-4o",
     model_source="auto",
+    unique_model=None,
+    unique_model_source="auto",
+    unique_passes=1,
+    merge_model=None,
+    merge_model_source="auto",
     creativity=0,
     max_workers=1,
     random_state=None,
@@ -296,7 +330,10 @@ def collapse_themes(
         description (str): Data/question context, injected into the prompt — e.g.
             the survey question the categories came from. Helps the model judge
             which distinctions matter.
-        passes (int): Number of collapse iterations to run. Default 1.
+        passes (int | str): Number of collapse iterations, or "auto" to iterate
+            until the deterministic quality benchmark peaks (the recommended mode
+            for a final taxonomy — pair with aggressive=True). Default 1.
+        max_passes (int): Cap on iterations when passes="auto". Default 10.
         batch_size (int): Themes per LLM chunk (ceil(n / batch_size) calls per
             pass). Default 40.
         aggressive (bool): If True, use the conceptual-merge prompt (compress);
@@ -308,9 +345,28 @@ def collapse_themes(
             0.92. None or >=1.0 skips embeddings.
         shuffle (bool): Randomize order each pass so batch composition varies.
             Default True (improves convergence stability).
-        user_model (str): Model name. Default "gpt-4o". Use a capable model —
-            small models can degenerate into repetition.
+        final_consolidation (float): Cosine threshold for one greedy embedding
+            re-merge over the whole result after all passes, collapsing cross-batch
+            lexical-sibling duplicates that batched passes (and the auto loop) cannot
+            reach. Default 0.82 — deterministic and tuned to land just above the true
+            concept count (errs toward keeping categories; over-segmentation is
+            preferred over over-consolidation). False/None skips.
+        user_model (str): Model name for the merge phase. Default "gpt-4o". Use a
+            capable model — small models can degenerate into repetition.
         model_source (str): Provider — "auto", "openai", "huggingface", etc.
+        unique_model (str): If set, run an initial extract-unique thinning phase on
+            this (typically cheaper) model before the merge phase, allocating model
+            spend by task difficulty: a smaller model handles faithful restatement
+            removal, a stronger one handles conceptual merging. None (default) skips
+            the phase entirely (backward compatible). Recommended pairing:
+            unique_model = a 72B-class model, merge_model = a frontier model.
+        unique_model_source (str): Provider for unique_model. Default "auto" — can
+            differ from the merge phase, so the two phases may sit on different
+            providers.
+        unique_passes (int): Number of extract-unique passes in the thinning phase
+            when unique_model is set. Default 1.
+        merge_model (str): Model for the merge phase. Defaults to user_model when None.
+        merge_model_source (str): Provider for merge_model. Default "auto".
         creativity (float): Temperature. Default 0 (deterministic).
         max_workers (int): Batches processed concurrently per pass. Default 1.
         random_state (int): Seed for shuffling (per-pass seed = random_state + p).
@@ -325,37 +381,96 @@ def collapse_themes(
         >>> import cat_stack as cat
         >>> themes = cat.explore(df['responses'], description="Why did you move?",
         ...                      api_key=key)
-        >>> # 1) thin faithfully, then 2) compress
-        >>> thinned = cat.collapse_themes(themes, api_key=key,
-        ...     description="Why did you move?", passes=10, max_workers=8)
-        >>> final = cat.collapse_themes(thinned, api_key=key,
-        ...     description="Why did you move?", passes=2, aggressive=True)
+        >>> # Recommended: aggressive merge, auto-stop at the quality peak
+        >>> taxonomy = cat.collapse_themes(
+        ...     themes, api_key=key, description="Why did you move?",
+        ...     aggressive=True, passes="auto", max_workers=8,
+        ... )
     """
     if not api_key:
         raise ValueError("collapse_themes() needs an api_key for the LLM call.")
 
     mode = "merge" if aggressive else "unique"
-    provider = detect_provider(user_model, model_source)
-    client = UnifiedLLMClient(provider=provider, api_key=api_key, model=user_model)
 
-    current = input_data
-    for p in range(passes):
-        seed = None if random_state is None else random_state + p
-        current = _collapse_once(
-            client,
-            current,
+    # The main (merge) phase runs on merge_model if given, else user_model. A separate
+    # cheaper model can handle the simpler unique-keeping phase via unique_model — per
+    # step the work differs in difficulty (faithful thinning is easy, conceptual
+    # merging is hard), so model spend can be allocated accordingly. Each phase
+    # resolves its own provider, so the two can sit on different providers.
+    merge_name = merge_model or user_model
+    merge_src = merge_model_source if merge_model else model_source
+    merge_provider = detect_provider(merge_name, merge_src)
+    client = UnifiedLLMClient(provider=merge_provider, api_key=api_key, model=merge_name)
+
+    def _run(cl, items, md, p):
+        return _collapse_once(
+            cl, items,
             description=description,
             batch_size=batch_size,
             dedupe_threshold=dedupe_threshold,
             embedding_merge_threshold=embedding_merge_threshold,
-            mode=mode,
+            mode=md,
             shuffle=shuffle,
-            random_state=seed,
+            random_state=(None if random_state is None else random_state + p),
             creativity=creativity,
             max_workers=max_workers,
         )
-        if progress_callback:
-            progress_callback(p + 1, passes, "collapse_themes")
+
+    def _pass(items, p):
+        return _run(client, items, mode, p)
+
+    current = input_data
+
+    # Phase 1 (optional): cheap unique-keeping thin. When unique_model is set, run
+    # `unique_passes` extract-unique passes on a separate (typically smaller, cheaper)
+    # model to strip restatement-level duplicates before the expensive merge phase.
+    # Skipped entirely when unique_model is None (fully backward compatible).
+    if unique_model:
+        u_provider = detect_provider(unique_model, unique_model_source)
+        u_client = UnifiedLLMClient(provider=u_provider, api_key=api_key, model=unique_model)
+        for p in range(int(unique_passes)):
+            current = _run(u_client, current, "unique", p)
+            if progress_callback:
+                progress_callback(p + 1, int(unique_passes), "collapse_themes:unique")
+    if passes == "auto":
+        # Iterate until the deterministic quality benchmark stops improving (the
+        # peak), capped at max_passes. Quality is scored vs the ORIGINAL input
+        # themes — embedding-only, model-independent at decision time. The peak is
+        # the principled stop (validated across surveys and list sizes).
+        raw_embs = _get_emb_model().encode(
+            list(_to_counts(input_data).keys()), normalize_embeddings=True,
+            show_progress_bar=False,
+        )
+        best, best_q = None, -1.0
+        for p in range(max_passes):
+            current = _pass(current, p)
+            q = _quality(current, raw_embs)
+            if progress_callback:
+                progress_callback(p + 1, max_passes, "collapse_themes")
+            if q < best_q:
+                break  # quality dropped -> the previous pass was the peak
+            best, best_q = current, q
+        current = best if best is not None else current
+    else:
+        for p in range(int(passes)):
+            current = _pass(current, p)
+            if progress_callback:
+                progress_callback(p + 1, int(passes), "collapse_themes")
+
+    # Final global consolidation. Batched passes (and the auto loop) can only merge
+    # labels that share a batch, so cross-batch lexical siblings — e.g. "tension" vs
+    # "estrangement", which restate one concept but embed below the per-pass dedupe
+    # threshold — survive as separate themes, inflating the count above the true
+    # number of concepts. This applies one greedy embedding re-merge over the WHOLE
+    # result at a lower threshold, dropping each label that restates an already-kept
+    # one to bring the count closer to truth. Greedy (compares only against kept
+    # representatives, no transitive chaining) avoids blobbing related-but-distinct
+    # labels. It is deterministic (no extra LLM call, model-independent at decision
+    # time) and tuned to land just above the true count, so it errs toward KEEPING
+    # categories — over-segmentation is the preferred failure mode, not
+    # over-consolidation. Set final_consolidation=False to skip.
+    if final_consolidation and len(current) > 1:
+        current = _embedding_merge(current, final_consolidation)
 
     if filename:
         pd.DataFrame({"category": current}).to_csv(filename, index=False)