cat-stack 1.6.8__tar.gz → 2.0.0b1__tar.gz

{cat_stack-1.6.8 → cat_stack-2.0.0b1}/.gitignore

@@ -41,3 +41,11 @@ survey-summarizer
 # R package build/check artifacts
 *.Rcheck/
 *.tar.gz
+
+# Local experiment scratch (not part of the package)
+/patches/
+/AGENTS.md
+/test_live_*.py
+/test_smoke_*.py
+/test_stress_*.py
+/test_parallel_*.py

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

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

{cat_stack-1.6.8 → cat_stack-2.0.0b1}/src/catstack/__init__.py

@@ -18,6 +18,7 @@ from .__about__ import (
 # Main entry points
 from .extract import extract
 from .explore import explore
+from .collapse_themes import collapse_themes
 from .classify import classify
 from .summarize import summarize
 from .prompt_tune import prompt_tune
@@ -103,6 +104,7 @@ __all__ = [
     # Main entry points
     "extract",
     "explore",
+    "collapse_themes",
     "classify",
     "summarize",
     "prompt_tune",

{cat_stack-1.6.8 → cat_stack-2.0.0b1}/src/catstack/_formatter.py

@@ -45,6 +45,11 @@ def _check_dependencies():
 def _check_dependencies_installed() -> bool:
     """Pure check — returns True if all formatter deps import successfully.
     No side effects, no install attempts."""
+    # If a dep was just pip-installed in this process's lifetime, the import
+    # system may have cached its earlier absence; clear that so the re-check
+    # actually sees the freshly-installed package.
+    import importlib
+    importlib.invalidate_caches()
     try:
         import torch  # noqa: F401
         import transformers  # noqa: F401
@@ -165,7 +170,31 @@ def _ensure_dependencies(verbose: bool = True) -> bool:
             "  To skip this and disable the formatter, pass json_formatter=False."
         )
 
-    return _install_dependencies(verbose=verbose)
+    ok = _install_dependencies(verbose=verbose)
+    if not ok:
+        # Freshly pip-installed packages (esp. compiled ones like torch) often
+        # cannot be imported by the SAME running process — but they ARE on disk
+        # now. Tell the user to re-run rather than silently degrading every row
+        # to an error.
+        if verbose and _deps_on_disk():
+            print(
+                "[CatLLM] Formatter dependencies were just installed but cannot "
+                "be imported into the already-running process. Please RE-RUN your "
+                "command — they will load on the next start. (Avoid this by "
+                "pre-installing: pip install 'cat-stack[formatter]'.)"
+            )
+    return ok
+
+
+def _deps_on_disk() -> bool:
+    """True if the formatter deps are findable on disk (importable in a FRESH
+    process) even if they failed to import in the current one."""
+    import importlib.util
+    try:
+        return all(importlib.util.find_spec(m) is not None
+                   for m in ("torch", "transformers", "accelerate"))
+    except (ImportError, ValueError):
+        return False
 
 
 def _is_model_cached() -> bool:
@@ -205,6 +234,51 @@ def ensure_formatter_available() -> bool:
     return True  # actual download happens in load_formatter()
 
 
+def _load_formatter_tokenizer(AutoTokenizer):
+    """Load the formatter tokenizer, defending against a malformed
+    `tokenizer_config.json`.
+
+    Some published configs store `extra_special_tokens` as a LIST, but
+    transformers 4.56–4.57.x expect a dict and crash in
+    `_set_model_specific_special_tokens` with
+    `'list' object has no attribute 'keys'`. On that failure we snapshot the
+    repo locally, normalize a list-valued `extra_special_tokens` to `{}`
+    (the tokens already live in `added_tokens`/`special_tokens_map`, so
+    dropping the field is lossless), and load from the patched local copy.
+    """
+    try:
+        return AutoTokenizer.from_pretrained(
+            _MERGED_MODEL_REPO, trust_remote_code=True
+        )
+    except (AttributeError, TypeError) as e:
+        if "keys" not in str(e) and "extra_special_tokens" not in str(e):
+            raise
+        import json
+        import os
+        from huggingface_hub import snapshot_download
+
+        local_dir = snapshot_download(_MERGED_MODEL_REPO)
+        cfg_path = os.path.join(local_dir, "tokenizer_config.json")
+        with open(cfg_path) as f:
+            cfg = json.load(f)
+        if isinstance(cfg.get("extra_special_tokens"), list):
+            cfg["extra_special_tokens"] = {}
+            # snapshot dirs are often read-only symlink caches; patch a copy.
+            import tempfile
+            import shutil
+            patched = tempfile.mkdtemp(prefix="catllm_formatter_tok_")
+            for fn in os.listdir(local_dir):
+                src = os.path.join(local_dir, fn)
+                if os.path.isfile(src):
+                    shutil.copy(src, os.path.join(patched, fn))
+            with open(os.path.join(patched, "tokenizer_config.json"), "w") as f:
+                json.dump(cfg, f)
+            print("[CatLLM] Patched malformed extra_special_tokens in the "
+                  "formatter tokenizer config (list -> {}).")
+            return AutoTokenizer.from_pretrained(patched, trust_remote_code=True)
+        raise
+
+
 def load_formatter(device=None):
     """
     Load the merged formatter model and tokenizer.
@@ -230,15 +304,21 @@ def load_formatter(device=None):
     dtype = torch.float16 if device == "cuda" else torch.float32
 
     print(f"[CatLLM] Loading JSON formatter on {device}...")
-    tokenizer = AutoTokenizer.from_pretrained(
-        _MERGED_MODEL_REPO, trust_remote_code=True
-    )
+    tokenizer = _load_formatter_tokenizer(AutoTokenizer)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    model = AutoModelForCausalLM.from_pretrained(
-        _MERGED_MODEL_REPO, dtype=dtype, trust_remote_code=True
-    )
+    # `dtype=` is the transformers >=4.56 kwarg; older versions only accept
+    # `torch_dtype=` and crash if `dtype=` leaks into the config. Try the new
+    # name, fall back to the old one.
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            _MERGED_MODEL_REPO, dtype=dtype, trust_remote_code=True
+        )
+    except TypeError:
+        model = AutoModelForCausalLM.from_pretrained(
+            _MERGED_MODEL_REPO, torch_dtype=dtype, trust_remote_code=True
+        )
     model = model.to(device)
     model.eval()
 
@@ -281,7 +361,9 @@ def run_formatter(raw_output, categories, model, tokenizer, device):
     with torch.no_grad():
         out = model.generate(
             **inputs,
-            max_new_tokens=128,
+            # 512 (was 128): a large category set produces a long N-key JSON
+            # object; 128 tokens truncated it for 28/48-category tasks.
+            max_new_tokens=512,
             do_sample=False,
             temperature=None,
             top_p=None,

cat_stack-2.0.0b1/src/catstack/collapse_themes.py

@@ -0,0 +1,364 @@
+"""
+Theme collapsing for CatLLM.
+
+collapse_themes() takes an already-extracted list of category/theme strings (for
+example the output of explore()) and iteratively consolidates near-duplicate /
+synonymous labels into a smaller list. Each pass:
+
+    A. accept the list,
+    B. PRE-CLEAN before the model — normalize + Jaro-Winkler dedup (surface
+       variants) then embedding-merge (semantic near-duplicates),
+    C. split the cleaned list into batches of `batch_size`,
+    D. read every batch with one LLM call (extract-unique, or aggressive merge),
+    E. concatenate and dedupe into a single, smaller list.
+
+`passes` iterations run in one call, randomizing batch composition each pass so
+labels stranded in separate batches get fresh chances to meet and merge.
+Provider-agnostic via the same dispatch classify()/explore() use.
+"""
+
+import random
+import re
+import sys
+
+import numpy as np
+import pandas as pd
+from jellyfish import jaro_winkler_similarity
+
+from ._providers import UnifiedLLMClient, detect_provider
+from ._utils import _clean_label
+
+__all__ = [
+    "collapse_themes",
+]
+
+_LINE_PAT = re.compile(r"^\s*\d+\s*[\.\)\-]\s*(.+)$")
+_EMB_MODEL = None  # cached embedding model (loaded once per process)
+
+
+def _strip_parens(label):
+    """Drop parenthetical examples — '(...)' doesn't change the category."""
+    return re.sub(r"\s*\([^)]*\)", "", label).strip()
+
+
+def _norm_key(label):
+    """Canonical dedup key: parens-stripped, lowercased, separators/order unified."""
+    s = _strip_parens(label).lower().strip()
+    s = re.sub(r"\s*&\s*|\s+and\s+|\s*/\s*", " / ", s)
+    parts = sorted(p.strip() for p in s.split("/") if p.strip())
+    return " / ".join(parts)
+
+
+def _jw_dedupe(items, threshold):
+    """Order-preserving dedup: normalize each label and collapse near-identical
+    normalized labels with a Jaro-Winkler threshold. Returns readable forms."""
+    kept_keys = []
+    out = []
+    for c in items:
+        disp = _strip_parens(c).lower().strip()
+        key = _norm_key(c)
+        if not disp or not key:
+            continue
+        is_dup = any(
+            k == key
+            or (threshold < 1.0 and jaro_winkler_similarity(key, k) >= threshold)
+            for k in kept_keys
+        )
+        if not is_dup:
+            kept_keys.append(key)
+            out.append(disp)
+    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)."""
+    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)
+    reps, rep_embs = [], []
+    for it, e in zip(items, embs):
+        if rep_embs and float(np.max(np.asarray(rep_embs) @ e)) >= threshold:
+            continue
+        reps.append(it)
+        rep_embs.append(e)
+    return reps
+
+
+def _collapse_batch(client, batch, description, creativity, mode="unique"):
+    """One LLM call on a single batch -> list[str].
+
+    mode="unique": extract unique categories only (remove restatements, keep
+    distinct ones) — gentle, near-idempotent, guaranteed to only remove.
+    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).
+    """
+    items_blob = "; ".join(batch)
+    context = f' about: "{description}"' if description else ""
+    if mode == "merge":
+        prompt = (
+            f"You are consolidating a list of category labels{context} into a smaller set of "
+            "broader categories. Group labels that describe the same underlying concept and give "
+            "each group ONE clear representative label — actively merge near-synonyms and closely "
+            "related labels into broader themes. BUT retain nuance: do NOT over-merge — keep labels "
+            "separate when they capture a genuinely distinct concept, even if related, rather than "
+            "collapsing them into one vague catch-all. Prefer fewer, cleaner categories without "
+            f"losing real distinctions. Labels are separated by semicolons within triple backticks: "
+            f"```{items_blob}```.\n\n"
+            "Return ONLY a numbered list of the consolidated categories. 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}. "
+            "Return the UNIQUE categories. Remove ONLY exact duplicates and labels that "
+            "restate the SAME category in different words — when two labels are the same "
+            "category, keep one of them exactly as written. KEEP every genuinely distinct "
+            "category. Do NOT merge categories that are merely related, do NOT invent or "
+            "broaden labels, and do NOT drop a category just to make the list shorter. "
+            "If all the labels are already distinct categories, return ALL of them 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"
+        )
+    reply, error = client.complete(
+        messages=[{"role": "user", "content": prompt}],
+        creativity=creativity,
+        force_json=False,
+    )
+    if error:
+        # No data loss: keep the batch unchanged so its categories aren't dropped.
+        sys.stderr.write(f"[collapse_themes] batch failed: {error} — keeping batch unchanged\n")
+        return [str(x).strip().lower() for x in batch]
+
+    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)
+
+    if mode == "unique":
+        # Contraction guarantee: extract-unique 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.
+        in_by_key = {}
+        for x in batch:
+            in_by_key.setdefault(_norm_key(x), str(x).strip().lower())
+        seen, subset = set(), []
+        for o in out:
+            k = _norm_key(o)
+            if k in in_by_key and k not in seen:
+                seen.add(k)
+                subset.append(in_by_key[k])
+        # If parsing/matching failed entirely, fall back to the batch (no loss).
+        out = subset if subset else [str(x).strip().lower() for x in batch]
+    return out
+
+
+def _to_counts(input_data):
+    """Coerce the accepted input forms into a {category: count} dict."""
+    if isinstance(input_data, pd.DataFrame):
+        cols = {c.lower(): c for c in input_data.columns}
+        cat_col = cols.get("category")
+        cnt_col = cols.get("count")
+        if cat_col is None:
+            raise ValueError("DataFrame input must have a 'category' column.")
+        if cnt_col is not None:
+            return input_data.groupby(cat_col)[cnt_col].sum().astype(int).to_dict()
+        return input_data[cat_col].value_counts().to_dict()
+    if isinstance(input_data, dict):
+        return {str(k): int(v) for k, v in input_data.items()}
+    series = input_data if isinstance(input_data, pd.Series) else pd.Series(input_data)
+    series = series.dropna().astype("string")
+    return series.value_counts().to_dict()
+
+
+def _collapse_once(
+    client,
+    items,
+    *,
+    description,
+    batch_size,
+    dedupe_threshold,
+    embedding_merge_threshold,
+    mode,
+    shuffle,
+    random_state,
+    creativity,
+    max_workers,
+):
+    """Run a single collapse pass over `items` and return the reduced list."""
+    # A. accept -> {category: count}
+    counts = _to_counts(items)
+
+    # B. PRE-CLEAN before the model: normalize+JW dedup, then embedding-merge
+    ordered = sorted(counts, key=counts.get, reverse=True)
+    cleaned = _jw_dedupe(ordered, dedupe_threshold)
+    cleaned = _embedding_merge(cleaned, embedding_merge_threshold)
+
+    # Randomize order so batch composition varies across passes — gives near-
+    # duplicates split across batches fresh chances to co-occur and merge.
+    if shuffle:
+        random.Random(random_state).shuffle(cleaned)
+
+    # C. split into batches
+    batches = [cleaned[i:i + batch_size] for i in range(0, len(cleaned), batch_size)]
+
+    # D. one LLM call per batch (sequential or parallel)
+    if max_workers and max_workers > 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, mode): 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 batch in batches:
+            out.extend(_collapse_batch(client, batch, description, creativity, mode))
+
+    # E. dedupe the concatenated output (surface-level)
+    return _jw_dedupe(out, dedupe_threshold)
+
+
+def collapse_themes(
+    input_data,
+    api_key=None,
+    description="",
+    passes=1,
+    batch_size=40,
+    aggressive=False,
+    dedupe_threshold=0.95,
+    embedding_merge_threshold=0.92,
+    shuffle=True,
+    user_model="gpt-4o",
+    model_source="auto",
+    creativity=0,
+    max_workers=1,
+    random_state=None,
+    filename=None,
+    progress_callback=None,
+):
+    """
+    Collapse a list of extracted themes into a smaller, deduplicated list.
+
+    Iteratively consolidates near-duplicate / synonymous category labels (for
+    example the output of explore()). Each pass PRE-CLEANS before the model
+    (normalize + Jaro-Winkler dedup, then embedding-merge), splits into batches,
+    sends each batch to the model, and dedupes the concatenated result. Runs
+    `passes` iterations, randomizing batch composition each pass so labels
+    stranded in separate batches get fresh chances to merge.
+
+    Two modes:
+      - aggressive=False (default): extract-unique — only removes duplicates /
+        restatements, never invents or broadens. Each pass is guaranteed monotone
+        (output is a subset of its input). Use to thin a noisy list faithfully.
+      - aggressive=True: conceptual merge — actively consolidates related labels
+        into broader categories while retaining meaningful distinctions. Use as a
+        final compression step.
+
+    Provider-agnostic (model_source: "auto", "openai", "huggingface", ...), via
+    the same dispatch classify()/explore() use.
+
+    Args:
+        input_data: Themes to collapse. list[str] (duplicates allowed), pandas
+            Series, dict {category: count}, or DataFrame with "category"
+            [and optional "count"] columns.
+        api_key (str): API key for the model provider.
+        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.
+        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);
+            if False, extract-unique (faithful thinning). Default False.
+        dedupe_threshold (float): Jaro-Winkler similarity at/above which two
+            normalized labels are deduped. Default 0.95; 1.0 = exact only.
+        embedding_merge_threshold (float): Cosine similarity at/above which labels
+            are merged in the pre-LLM embedding step (BAAI/bge-small). Default
+            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.
+        model_source (str): Provider — "auto", "openai", "huggingface", etc.
+        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).
+            None = nondeterministic.
+        filename (str): Optional CSV path to save the final list.
+        progress_callback (callable): Optional callback(pass, passes, label).
+
+    Returns:
+        list[str]: The collapsed category list after `passes` iterations.
+
+    Examples:
+        >>> 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)
+    """
+    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,
+            description=description,
+            batch_size=batch_size,
+            dedupe_threshold=dedupe_threshold,
+            embedding_merge_threshold=embedding_merge_threshold,
+            mode=mode,
+            shuffle=shuffle,
+            random_state=seed,
+            creativity=creativity,
+            max_workers=max_workers,
+        )
+        if progress_callback:
+            progress_callback(p + 1, passes, "collapse_themes")
+
+    if filename:
+        pd.DataFrame({"category": current}).to_csv(filename, index=False)
+        print(f"Collapsed categories saved to {filename}")
+
+    return current