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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cat-stack
-Version: 2.0.0b1
+Version: 2.0.0b4
 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

{cat_stack-2.0.0b1 → cat_stack-2.0.0b4}/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.0b4"
 __author__ = "Chris Soria"
 __email__ = "chrissoria@berkeley.edu"
 __title__ = "cat-stack"

{cat_stack-2.0.0b1 → cat_stack-2.0.0b4}/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)