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

{cat_stack-2.4.0 → cat_stack-2.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cat-stack
-Version: 2.4.0
+Version: 2.5.0
 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.4.0 → cat_stack-2.5.0}/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.4.0"
+__version__ = "2.5.0"
 __author__ = "Chris Soria"
 __email__ = "chrissoria@berkeley.edu"
 __title__ = "cat-stack"

{cat_stack-2.4.0 → cat_stack-2.5.0}/src/catstack/collapse_themes.py

@@ -122,13 +122,18 @@ def _collapse_batch(client, batch, description, creativity, mode="unique"):
 
     mode="unique": extract unique categories only (remove restatements, keep
     distinct ones) — gentle, near-idempotent, guaranteed to only remove.
+    mode="prune": drop CONCEPTUAL duplicates (labels naming the same concept in
+    different words), keeping one representative per concept VERBATIM — never
+    renames or merges merely-related labels. Stronger than "unique" (catches
+    same-concept-different-words) but, unlike "merge", never collapses distinct
+    concepts. Output forced to a subset of the input.
     mode="merge": aggressively consolidate related labels into broader concepts
     while retaining meaningful distinctions — for a final compression step.
 
     Strict numbered-list prompt + strict parsing, so the reply is always a clean
     list and any stray prose is ignored. Guardrails: a failed call returns the
-    batch unchanged (no data loss); in "unique" mode the output is forced to be a
-    subset of the input (monotone, drift-free).
+    batch unchanged (no data loss); in "unique"/"prune" modes the output is forced
+    to be a subset of the input (monotone, drift-free).
     """
     items_blob = "; ".join(batch)
     context = f' about: "{description}"' if description else ""
@@ -150,6 +155,23 @@ def _collapse_batch(client, batch, description, creativity, mode="unique"):
             "2. Education\n"
             "3. Religion"
         )
+    elif mode == "prune":
+        prompt = (
+            f"You are given a list of category labels{context}. Some labels are DUPLICATES — they "
+            "name the SAME concept as another label using different words. Return the list with "
+            "duplicates removed. Rules: keep ONE label per distinct concept, copied EXACTLY as "
+            "written (verbatim); do NOT rename, rephrase, merge, or broaden any label; keep two "
+            "labels SEPARATE whenever they name genuinely different concepts, even if related; if "
+            "no labels duplicate each other, return them ALL unchanged. "
+            f"Labels are separated by semicolons within triple backticks: ```{items_blob}```.\n\n"
+            "Return ONLY a numbered list, using the labels exactly as they appear. Each line must "
+            "follow this exact format, with no other text before or after the list:\n"
+            "N. label\n\n"
+            "Example:\n"
+            "1. Employment\n"
+            "2. Education\n"
+            "3. Religion"
+        )
     else:
         prompt = (
             f"You are given a list of category labels{context}. "
@@ -186,8 +208,16 @@ def _collapse_batch(client, batch, description, creativity, mode="unique"):
             if label:
                 out.append(label)
 
-    if mode == "unique":
-        # Contraction guarantee: extract-unique must only REMOVE, never add or
+    if not out:
+        # Unparseable reply (no numbered list found): keep the batch unchanged
+        # rather than silently dropping every label in it. Mirrors the API-error
+        # fallback above; matters most in "merge" mode, which has no
+        # subset-of-input guarantee to fall back on.
+        sys.stderr.write("[collapse_themes] unparseable reply — keeping batch unchanged\n")
+        return [str(x).strip().lower() for x in batch]
+
+    if mode in ("unique", "prune"):
+        # Contraction guarantee: extract-unique/prune must only REMOVE, never add or
         # mutate. Keep only outputs that map back to an input label (by normalized
         # key), as the original input string. Makes every pass monotone and
         # drift-free, immune to intermittent model rephrasing/splitting.
@@ -205,6 +235,72 @@ def _collapse_batch(client, batch, description, creativity, mode="unique"):
     return out
 
 
+def _select_top_n(client, counts, n, description, creativity, dedupe_threshold):
+    """The "nuclear" final step: one global LLM call that consolidates the whole
+    surviving list into at most `n` categories, guided by generation counts
+    (higher = the theme recurred more often during extraction). Overlapping
+    labels are merged rather than dropped; only rare, unabsorbable themes fall
+    away. The result is guaranteed <= n: over-length replies are truncated, and
+    a failed or unparseable call falls back to the top-n labels by count
+    (deterministic)."""
+    ordered = sorted(counts, key=counts.get, reverse=True)
+    fallback = [str(x).strip().lower() for x in ordered[:n]]
+    blob = "; ".join(f"{lbl} ({counts[lbl]})" for lbl in ordered)
+    context = f' about: "{description}"' if description else ""
+    prompt = (
+        f"You are given category labels{context}. Each label is followed by a count in "
+        "parentheses: how many times it was independently generated during extraction — "
+        "higher counts mean the theme is more common in the data. Consolidate this list "
+        f"into EXACTLY {n} final categories that best summarize the data. Favor "
+        "high-count themes; merge overlapping or related labels into one clear "
+        "representative label rather than dropping them; drop only rare themes that "
+        f"cannot reasonably fold into any of the {n}. Labels are separated by semicolons "
+        f"within triple backticks: ```{blob}```\n\n"
+        f"Return ONLY a numbered list of exactly {n} category labels, without counts. "
+        "Each line must follow this exact format, with no other text before or after "
+        "the list:\n"
+        "N. label\n\n"
+        "Example:\n"
+        "1. Employment\n"
+        "2. Education\n"
+        "3. Religion"
+    )
+    reply, error = client.complete(
+        messages=[{"role": "user", "content": prompt}],
+        creativity=creativity,
+        force_json=False,
+    )
+    if error:
+        sys.stderr.write(f"[collapse_themes] top_n call failed: {error} — "
+                         "falling back to top-n by count\n")
+        return fallback
+    out = []
+    for line in (reply or "").splitlines():
+        m = _LINE_PAT.match(line.strip())
+        if m:
+            label = _clean_label(m.group(1)).strip(" ;.,")
+            if label:
+                out.append(label)
+    out = _jw_dedupe(out, dedupe_threshold)
+    if not out:
+        sys.stderr.write("[collapse_themes] top_n reply unparseable — "
+                         "falling back to top-n by count\n")
+        return fallback
+    return out[:n]
+
+
+def _count_guidance(current, input_data):
+    """{label: count} evidence for the top_n prompt: each surviving label gets its
+    aggregate generation count from the ORIGINAL input (matched by normalized key)
+    when available, else its count in the current list."""
+    cur = _to_counts(current)
+    orig = {}
+    for k, v in _to_counts(input_data).items():
+        key = _norm_key(k)
+        orig[key] = orig.get(key, 0) + int(v)
+    return {lbl: max(int(c), orig.get(_norm_key(lbl), 0)) for lbl, c in cur.items()}
+
+
 def _to_counts(input_data):
     """Coerce the accepted input forms into a {category: count} dict."""
     if isinstance(input_data, pd.DataFrame):
@@ -288,6 +384,9 @@ def collapse_themes(
     embedding_merge_threshold=0.92,
     shuffle=True,
     final_consolidation=0.82,
+    top_n=None,
+    prune=False,
+    prune_threshold=50,
     user_model="gpt-4o",
     model_source="auto",
     unique_model=None,
@@ -332,7 +431,9 @@ def collapse_themes(
             which distinctions matter.
         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.
+            for a final taxonomy — pair with aggressive=True). Either way,
+            iteration stops early once two consecutive passes leave the list
+            unchanged. 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.
@@ -345,12 +446,29 @@ 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).
+        prune (bool): If True, use the prune strategy instead of the merge/unique
+            passes: drop only CONCEPTUAL duplicates (one representative per concept,
+            kept verbatim), never renaming or merging merely-related labels, so a
+            short clean list is left intact while a long list is deduplicated. Lists
+            at/below prune_threshold go straight to a single global prune; longer
+            lists are reduced by batched prune first, then a final global prune.
+        prune_threshold (int): Max list length that goes directly to a single global
+            prune call (no batching). Default 50.
         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.
+        top_n (int): Optional "nuclear" final step. After all passes and the final
+            consolidation, ONE global LLM call consolidates the surviving list into
+            at most `top_n` categories, guided by each label's generation count
+            (frequent themes favored; overlapping labels merged into one
+            representative rather than dropped). None (default) changes nothing
+            about existing behavior. The result is guaranteed to have <= top_n
+            labels: a failed or unparseable call falls back deterministically to
+            the top_n labels by count. Use when a compact fixed-size taxonomy is
+            required and coverage of rare themes is knowingly traded away.
         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.
@@ -378,7 +496,7 @@ def collapse_themes(
         list[str]: The collapsed category list after `passes` iterations.
 
     Examples:
-        >>> import cat_stack as cat
+        >>> import catstack as cat
         >>> themes = cat.explore(df['responses'], description="Why did you move?",
         ...                      api_key=key)
         >>> # Recommended: aggressive merge, auto-stop at the quality peak
@@ -419,6 +537,63 @@ def collapse_themes(
     def _pass(items, p):
         return _run(client, items, mode, p)
 
+    # ── Prune strategy ───────────────────────────────────────────────────────
+    # Drop CONCEPTUAL duplicates only (keep one representative per concept,
+    # verbatim) — never rename or merge merely-related labels, so it cannot
+    # over-consolidate a short, already-clean list. Length-routed: a list at or
+    # below prune_threshold goes straight to a single global "master list" prune;
+    # a longer list is first reduced by the same prune primitive in shuffled
+    # batches (until it stops shrinking or fits the threshold), then finished with
+    # the global master prune.
+    if prune:
+        items = list(_to_counts(input_data).keys())
+        items = _jw_dedupe(items, dedupe_threshold)              # gentle, no rename
+        items = _embedding_merge(items, embedding_merge_threshold)
+        rng = random.Random(random_state)
+        guard = 0
+        while len(items) > prune_threshold and guard < max_passes:
+            guard += 1
+            if shuffle:
+                rng.shuffle(items)
+            batches = [items[i:i + prune_threshold]
+                       for i in range(0, len(items), prune_threshold)]
+            if max_workers and max_workers > 1 and len(batches) > 1:
+                from concurrent.futures import ThreadPoolExecutor, as_completed
+
+                results = [None] * len(batches)
+                with ThreadPoolExecutor(max_workers=max_workers) as ex:
+                    futures = {
+                        ex.submit(_collapse_batch, client, b, description,
+                                  creativity, "prune"): i
+                        for i, b in enumerate(batches)
+                    }
+                    for fut in as_completed(futures):
+                        results[futures[fut]] = fut.result()
+                out = [label for r in results for label in (r or [])]
+            else:
+                out = []
+                for b in batches:
+                    out += _collapse_batch(client, b, description, creativity,
+                                           mode="prune")
+            out = _jw_dedupe(out, dedupe_threshold)
+            if progress_callback:
+                progress_callback(guard, max_passes, "collapse_themes:prune")
+            if len(out) >= len(items):     # no further reduction this round
+                items = out
+                break
+            items = out
+        if len(items) > 1:                  # final global master-list prune
+            items = _jw_dedupe(
+                _collapse_batch(client, items, description, creativity, mode="prune"),
+                dedupe_threshold)
+        if top_n and len(items) > int(top_n):
+            items = _select_top_n(client, _count_guidance(items, input_data),
+                                  int(top_n), description, creativity, dedupe_threshold)
+        if filename:
+            pd.DataFrame({"category": items}).to_csv(filename, index=False)
+            print(f"Collapsed categories saved to {filename}")
+        return items
+
     current = input_data
 
     # Phase 1 (optional): cheap unique-keeping thin. When unique_model is set, run
@@ -442,6 +617,7 @@ def collapse_themes(
             show_progress_bar=False,
         )
         best, best_q = None, -1.0
+        prev_key, stable = None, 0
         for p in range(max_passes):
             current = _pass(current, p)
             q = _quality(current, raw_embs)
@@ -450,12 +626,26 @@ def collapse_themes(
             if q < best_q:
                 break  # quality dropped -> the previous pass was the peak
             best, best_q = current, q
+            key = tuple(sorted(current))
+            stable = stable + 1 if key == prev_key else 0
+            prev_key = key
+            if stable >= 2:
+                break  # converged: two consecutive passes changed nothing
         current = best if best is not None else current
     else:
+        prev_key, stable = None, 0
         for p in range(int(passes)):
             current = _pass(current, p)
             if progress_callback:
                 progress_callback(p + 1, int(passes), "collapse_themes")
+            # Early stop at a fixed point: shuffling gives labels one fresh batch
+            # composition to merge under; if two consecutive passes both change
+            # nothing, further passes are near-certain no-ops — stop burning calls.
+            key = tuple(sorted(current))
+            stable = stable + 1 if key == prev_key else 0
+            prev_key = key
+            if stable >= 2:
+                break
 
     # 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
@@ -470,8 +660,20 @@ def collapse_themes(
     # categories — over-segmentation is the preferred failure mode, not
     # over-consolidation. Set final_consolidation=False to skip.
     if final_consolidation and len(current) > 1:
+        # Order by original generation frequency first, so the greedy merge keeps
+        # the most frequently generated variant as each cluster's representative
+        # instead of whichever label the last (shuffled) pass happened to emit
+        # first. Labels the merge phase renamed have no original count and sort
+        # last, preserving their relative order (sort is stable).
+        orig_counts = {_norm_key(k): v for k, v in _to_counts(input_data).items()}
+        current = sorted(current, key=lambda c: orig_counts.get(_norm_key(c), 0),
+                         reverse=True)
         current = _embedding_merge(current, final_consolidation)
 
+    if top_n and len(_to_counts(current)) > int(top_n):
+        current = _select_top_n(client, _count_guidance(current, input_data),
+                                int(top_n), description, creativity, dedupe_threshold)
+
     if filename:
         pd.DataFrame({"category": current}).to_csv(filename, index=False)
         print(f"Collapsed categories saved to {filename}")

{cat_stack-2.4.0 → cat_stack-2.5.0}/src/catstack/explore.py

@@ -77,7 +77,7 @@ def explore(
         every iteration. Length ≈ iterations × divisions × categories_per_chunk.
 
     Examples:
-        >>> import cat_stack as cat
+        >>> import catstack as cat
         >>>
         >>> raw_categories = cat.explore(
         ...     input_data=df['responses'],

{cat_stack-2.4.0 → cat_stack-2.5.0}/src/catstack/extract.py

@@ -36,6 +36,8 @@ from .pdf_functions import (
     explore_pdf_categories,
 )
 
+from .collapse_themes import collapse_themes
+
 
 def extract(
     input_data,
@@ -61,6 +63,9 @@ def extract(
     auto_download: bool = False,
     input_mode=None,
     domain: str = "neutral",
+    engine: str = "collapse",
+    max_workers: int = 1,
+    collapse_kwargs: dict = None,
 ):
     """
     Unified category extraction function for text, image, and PDF inputs.
@@ -111,15 +116,31 @@ def extract(
             limits. Default 0.0 (no delay).
         auto_download (bool): If True, automatically download missing Ollama
             models without prompting. Default False.
+        engine (str): Consolidation engine for text input. "collapse" (default):
+            run raw extraction (as explore() does) and consolidate the FULL
+            inventory with collapse_themes() — semantic pre-clean, quality-
+            controlled passes, then a count-guided reduction to at most
+            `max_categories`. "legacy": the pre-2.5 single merge call, which
+            truncates the inventory to the top max_categories*3 labels by
+            exact-string count before merging; kept for reproducing older runs.
+        max_workers (int): Parallel API calls for both the extraction chunks and
+            the consolidation batches (text engine="collapse" only; extraction
+            also honors it under "legacy"). Default 1.
+        collapse_kwargs (dict): Optional overrides forwarded to collapse_themes()
+            when engine="collapse" — e.g. {"prune": True} or
+            {"passes": 2, "aggressive": False}. Defaults applied first:
+            passes="auto", aggressive=True. `top_n` cannot be overridden here;
+            it is always max_categories.
 
     Returns:
         dict with keys:
             - counts_df: DataFrame of categories with counts
             - top_categories: List of top category names
-            - raw_top_text: Raw model output from final merge step
+            - raw_top_text: Raw model output from final merge step ("" when
+              engine="collapse", which has no single merge reply)
 
     Examples:
-        >>> import cat_stack as cat
+        >>> import catstack as cat
         >>>
         >>> # Extract categories from text responses
         >>> results = cat.extract(
@@ -171,7 +192,41 @@ def extract(
     resolved_description = description or ""
 
     if input_type == "text":
-        return explore_common_categories(
+        if engine not in ("collapse", "legacy"):
+            raise ValueError(f"engine must be 'collapse' or 'legacy', got '{engine}'")
+
+        if engine == "legacy":
+            return explore_common_categories(
+                input_data=input_data,
+                api_key=api_key,
+                survey_question=resolved_description,
+                max_categories=max_categories,
+                categories_per_chunk=categories_per_chunk,
+                divisions=divisions,
+                user_model=user_model,
+                creativity=creativity,
+                specificity=specificity,
+                research_question=research_question,
+                filename=filename,
+                model_source=model_source,
+                iterations=iterations,
+                random_state=random_state,
+                focus=focus,
+                progress_callback=progress_callback,
+                chunk_delay=chunk_delay,
+                auto_download=auto_download,
+                max_workers=max_workers,
+                domain=domain,
+            )
+
+        # engine="collapse": raw extraction (what explore() does), then consolidate
+        # the FULL inventory with collapse_themes(). Unlike the legacy merge, no
+        # label is truncated away before consolidation, pre-cleaning is semantic
+        # (Jaro-Winkler + embeddings) rather than exact-string, and the final
+        # count-guided top_n step guarantees at most max_categories categories.
+        import pandas as pd
+
+        raw_items = explore_common_categories(
             input_data=input_data,
             api_key=api_key,
             survey_question=resolved_description,
@@ -182,17 +237,63 @@ def extract(
             creativity=creativity,
             specificity=specificity,
             research_question=research_question,
-            filename=filename,
+            filename=None,
             model_source=model_source,
             iterations=iterations,
             random_state=random_state,
             focus=focus,
             progress_callback=progress_callback,
+            return_raw=True,
             chunk_delay=chunk_delay,
             auto_download=auto_download,
+            max_workers=max_workers,
             domain=domain,
         )
 
+        # Frequency inventory in the same shape the legacy engine returned.
+        def _normalize(cat):
+            return "/".join(sorted(t.strip().lower() for t in str(cat).split("/")))
+
+        flat = [str(x).strip() for x in raw_items if str(x).strip()]
+        if not flat:
+            raise ValueError("No categories were extracted from the model responses.")
+        inv = pd.DataFrame(flat, columns=["Category"])
+        inv["normalized"] = inv["Category"].map(_normalize)
+        counts_df = (
+            inv.groupby("normalized")
+               .agg(Category=("Category", lambda x: x.value_counts().index[0]),
+                    counts=("Category", "size"))
+               .sort_values("counts", ascending=False)
+               .reset_index(drop=True)
+        )
+
+        ck = dict(passes="auto", aggressive=True)
+        ck.update(collapse_kwargs or {})
+        ck["top_n"] = int(max_categories)  # the required N — not overridable
+        top = collapse_themes(
+            raw_items,
+            api_key=api_key,
+            description=resolved_description,
+            user_model=user_model,
+            model_source=model_source,
+            creativity=0 if creativity is None else creativity,
+            max_workers=max_workers,
+            random_state=random_state,
+            progress_callback=progress_callback,
+            **ck,
+        )
+
+        if filename:
+            pd.DataFrame({"rank": range(1, len(top) + 1), "category": top}).to_csv(
+                filename, index=False)
+            print(f"Top {len(top)} categories saved to {filename}")
+
+        return {
+            "counts_df": counts_df,
+            "top_categories": top,
+            "raw_top_text": "",
+        }
+
     elif input_type == "image":
         return explore_image_categories(
             image_input=input_data,