cat-stack 0.1.0__py3-none-any.whl

cat_stack/_chunked.py

@@ -0,0 +1,424 @@
+"""
+Chunked classification for CatLLM.
+
+When users have large category lists, this module splits them into smaller
+chunks, runs a separate LLM call per chunk with local 1..N numbering, and
+merges the results back into global numbering so downstream code
+(aggregate_results, build_output_dataframes) sees a single merged JSON dict.
+
+Each chunk automatically gets a temporary "Other" catch-all category appended
+(unless one is already present in the chunk). This gives the LLM an escape
+hatch for ambiguous responses, improving classification accuracy. The "Other"
+column is dropped before merging back to global keys, so the final output
+only contains the user's real categories.
+"""
+
+import json
+import math
+
+from .text_functions import (
+    build_json_schema,
+    extract_json,
+    validate_classification_json,
+    ollama_two_step_classify,
+)
+from ._category_analysis import has_other_category
+
+
+def run_chunked_classification(
+    *,
+    client,
+    cfg,
+    item,
+    categories,
+    categories_str,
+    example_json,
+    json_schema,
+    cove_original_task,
+    effective_creativity,
+    use_json_schema,
+    survey_question,
+    survey_question_context,
+    examples_text,
+    chain_of_thought,
+    context_prompt,
+    step_back_prompt,
+    stepback_insights,
+    chain_of_verification,
+    thinking_budget,
+    max_retries,
+    multi_label,
+    categories_per_call,
+    add_unified_other=False,
+    formatter_fallback_fn,
+    # Mode-specific
+    is_pdf_mode,
+    is_image_mode,
+    pdf_mode=None,
+    pdf_dpi=150,
+    input_description="",
+    # Prompt builders (passed in to avoid circular imports)
+    build_text_prompt_fn=None,
+    build_pdf_prompt_fn=None,
+    build_image_prompt_fn=None,
+    google_multimodal_fn=None,
+    prepare_page_data_fn=None,
+    prepare_image_data_fn=None,
+    build_cove_prompts_fn=None,
+    run_cove_fn=None,
+):
+    """
+    Run chunked classification for one item across category chunks.
+
+    Splits the full category list into chunks of `categories_per_call`,
+    runs one LLM call per chunk, and merges results with key remapping.
+
+    Returns:
+        tuple: (json_result_str, error) — same contract as a single LLM call
+    """
+    # Build chunks: list of (chunk_categories, global_offset)
+    chunks = []
+    for start in range(0, len(categories), categories_per_call):
+        chunk_cats = categories[start : start + categories_per_call]
+        chunks.append((chunk_cats, start))
+
+    merged_json = {}
+    chunk_other_values = []  # Track per-chunk "Other" values for unification
+
+    for chunk_cats, global_offset in chunks:
+        # Add temporary "Other" catch-all if the chunk doesn't already have one.
+        # This gives the LLM an escape hatch for ambiguous responses, improving
+        # accuracy. The "Other" key is dropped before merging to global keys.
+        added_other = False
+        num_real_cats = len(chunk_cats)
+        if not has_other_category(chunk_cats):
+            chunk_cats_for_call = list(chunk_cats) + ["Other"]
+            added_other = True
+        else:
+            chunk_cats_for_call = chunk_cats
+
+        # Build chunk-local prompt components (with "Other" if added)
+        chunk_categories_str = "\n".join(
+            f"{j+1}. {cat}" for j, cat in enumerate(chunk_cats_for_call)
+        )
+        chunk_example_json = json.dumps(
+            {str(j + 1): "0" for j in range(len(chunk_cats_for_call))}, indent=2
+        )
+        chunk_json_schema = (
+            build_json_schema(
+                chunk_cats_for_call,
+                include_additional_properties=(cfg["provider"] != "google"),
+            )
+            if use_json_schema
+            else None
+        )
+
+        # Rebuild CoVe task for this chunk if CoVe enabled
+        chunk_cove_task = ""
+        if chain_of_verification:
+            if multi_label:
+                cove_categorize = "into the following categories"
+                cove_json = 'Provide your answer in JSON format where the category number is the key and "1" if present, "0" if not.'
+            else:
+                cove_categorize = "into the single most appropriate category"
+                cove_json = 'Provide your answer in JSON format where the category number is the key. Assign "1" to the single best matching category and "0" to all others.'
+            chunk_cove_task = f"""{survey_question_context}
+Categorize text responses {cove_categorize}:
+{chunk_categories_str}
+{cove_json}"""
+
+        # Run one LLM call for this chunk (with "Other" included)
+        chunk_result, chunk_error = _run_single_chunk_call(
+            client=client,
+            cfg=cfg,
+            item=item,
+            chunk_cats=chunk_cats_for_call,
+            chunk_categories_str=chunk_categories_str,
+            chunk_json_schema=chunk_json_schema,
+            chunk_example_json=chunk_example_json,
+            chunk_cove_task=chunk_cove_task,
+            effective_creativity=effective_creativity,
+            survey_question=survey_question,
+            survey_question_context=survey_question_context,
+            examples_text=examples_text,
+            chain_of_thought=chain_of_thought,
+            context_prompt=context_prompt,
+            step_back_prompt=step_back_prompt,
+            stepback_insights=stepback_insights,
+            chain_of_verification=chain_of_verification,
+            thinking_budget=thinking_budget,
+            max_retries=max_retries,
+            multi_label=multi_label,
+            formatter_fallback_fn=formatter_fallback_fn,
+            is_pdf_mode=is_pdf_mode,
+            is_image_mode=is_image_mode,
+            pdf_mode=pdf_mode,
+            pdf_dpi=pdf_dpi,
+            input_description=input_description,
+            build_text_prompt_fn=build_text_prompt_fn,
+            build_pdf_prompt_fn=build_pdf_prompt_fn,
+            build_image_prompt_fn=build_image_prompt_fn,
+            google_multimodal_fn=google_multimodal_fn,
+            prepare_page_data_fn=prepare_page_data_fn,
+            prepare_image_data_fn=prepare_image_data_fn,
+            build_cove_prompts_fn=build_cove_prompts_fn,
+            run_cove_fn=run_cove_fn,
+        )
+
+        if chunk_error:
+            return (json.dumps(merged_json) if merged_json else '{"1":"e"}', chunk_error)
+
+        # Remap chunk-local keys (1..N) to global keys, dropping "Other"
+        try:
+            chunk_parsed = json.loads(chunk_result)
+        except (json.JSONDecodeError, TypeError):
+            return ('{"1":"e"}', f"Failed to parse chunk result: {chunk_result}")
+
+        # The "Other" key (if added) is the last one: str(num_real_cats + 1)
+        other_local_key = str(num_real_cats + 1) if added_other else None
+
+        for local_key, value in chunk_parsed.items():
+            # Capture the temporary "Other" value, don't merge it
+            if local_key == other_local_key:
+                try:
+                    chunk_other_values.append(int(value))
+                except (ValueError, TypeError):
+                    chunk_other_values.append(0)
+                continue
+            try:
+                global_key = str(global_offset + int(local_key))
+                merged_json[global_key] = value
+            except (ValueError, TypeError):
+                # Non-numeric key — skip (shouldn't happen with proper schemas)
+                pass
+
+    # Unified "Other": if all real categories are 0 but at least one chunk's
+    # "Other" fired, the response genuinely doesn't fit any category.
+    if add_unified_other:
+        real_sum = sum(
+            int(v) for v in merged_json.values()
+            if str(v).strip() in ("0", "1")
+        )
+        other_sum = sum(chunk_other_values)
+        unified_other = "1" if real_sum == 0 and other_sum > 0 else "0"
+        merged_json[str(len(categories) + 1)] = unified_other
+
+    return (json.dumps(merged_json), None)
+
+
+def _run_single_chunk_call(
+    *,
+    client,
+    cfg,
+    item,
+    chunk_cats,
+    chunk_categories_str,
+    chunk_json_schema,
+    chunk_example_json,
+    chunk_cove_task,
+    effective_creativity,
+    survey_question,
+    survey_question_context,
+    examples_text,
+    chain_of_thought,
+    context_prompt,
+    step_back_prompt,
+    stepback_insights,
+    chain_of_verification,
+    thinking_budget,
+    max_retries,
+    multi_label,
+    formatter_fallback_fn,
+    is_pdf_mode,
+    is_image_mode,
+    pdf_mode,
+    pdf_dpi,
+    input_description,
+    build_text_prompt_fn,
+    build_pdf_prompt_fn,
+    build_image_prompt_fn,
+    google_multimodal_fn,
+    prepare_page_data_fn,
+    prepare_image_data_fn,
+    build_cove_prompts_fn,
+    run_cove_fn,
+):
+    """
+    Run one LLM call for one chunk of categories on one item.
+
+    Returns:
+        tuple: (json_result_str, error)
+    """
+    thinking_providers = ("google", "openai", "anthropic", "huggingface", "huggingface-together")
+
+    # =================================================================
+    # PDF MODE
+    # =================================================================
+    if is_pdf_mode and isinstance(item, tuple):
+        pdf_path, page_index, page_label = item
+
+        page_data = prepare_page_data_fn(
+            pdf_path=pdf_path,
+            page_index=page_index,
+            page_label=page_label,
+            pdf_mode=pdf_mode,
+            provider=cfg["provider"],
+            pdf_dpi=pdf_dpi,
+        )
+
+        if page_data.get("error"):
+            return ('{"1":"e"}', page_data["error"])
+
+        messages = build_pdf_prompt_fn(
+            page_data=page_data,
+            categories_str=chunk_categories_str,
+            input_description=input_description,
+            provider=cfg["provider"],
+            pdf_mode=pdf_mode,
+            chain_of_thought=chain_of_thought,
+            context_prompt=context_prompt,
+            step_back_prompt=step_back_prompt,
+            stepback_insights=stepback_insights,
+            model_name=cfg["model"],
+            example_json=chunk_example_json,
+            multi_label=multi_label,
+        )
+
+        if cfg["provider"] == "google":
+            reply, error = google_multimodal_fn(
+                client=client,
+                messages=messages,
+                json_schema=chunk_json_schema,
+                creativity=effective_creativity,
+                thinking_budget=thinking_budget,
+                max_retries=max_retries,
+            )
+        else:
+            reply, error = client.complete(
+                messages=messages,
+                json_schema=chunk_json_schema,
+                creativity=effective_creativity,
+                thinking_budget=thinking_budget if cfg["provider"] in thinking_providers else None,
+                max_retries=max_retries,
+            )
+
+        if error:
+            return ('{"1":"e"}', error)
+
+        json_result = extract_json(reply)
+        json_result = formatter_fallback_fn(json_result, reply, chunk_cats)
+        return (json_result, None)
+
+    # =================================================================
+    # IMAGE MODE
+    # =================================================================
+    elif is_image_mode and isinstance(item, tuple):
+        image_path, image_label = item
+
+        image_data = prepare_image_data_fn(image_path, image_label)
+
+        if image_data.get("error"):
+            return ('{"1":"e"}', image_data["error"])
+
+        messages = build_image_prompt_fn(
+            image_data=image_data,
+            categories_str=chunk_categories_str,
+            input_description=input_description,
+            provider=cfg["provider"],
+            chain_of_thought=chain_of_thought,
+            context_prompt=context_prompt,
+            step_back_prompt=step_back_prompt,
+            stepback_insights=stepback_insights,
+            model_name=cfg["model"],
+            example_json=chunk_example_json,
+            multi_label=multi_label,
+        )
+
+        if cfg["provider"] == "google":
+            reply, error = google_multimodal_fn(
+                client=client,
+                messages=messages,
+                json_schema=chunk_json_schema,
+                creativity=effective_creativity,
+                thinking_budget=thinking_budget,
+                max_retries=max_retries,
+            )
+        else:
+            reply, error = client.complete(
+                messages=messages,
+                json_schema=chunk_json_schema,
+                creativity=effective_creativity,
+                thinking_budget=thinking_budget if cfg["provider"] in thinking_providers else None,
+                max_retries=max_retries,
+            )
+
+        if error:
+            return ('{"1":"e"}', error)
+
+        json_result = extract_json(reply)
+        json_result = formatter_fallback_fn(json_result, reply, chunk_cats)
+        return (json_result, None)
+
+    # =================================================================
+    # TEXT MODE
+    # =================================================================
+    else:
+        response_text = item
+
+        if cfg["use_two_step"]:  # Ollama
+            json_result, error = ollama_two_step_classify(
+                client=client,
+                response_text=response_text,
+                categories=chunk_cats,
+                categories_str=chunk_categories_str,
+                survey_question=survey_question,
+                creativity=effective_creativity,
+                max_retries=max_retries,
+            )
+            if not error:
+                json_result = formatter_fallback_fn(json_result, json_result, chunk_cats)
+            return (json_result, error)
+        else:
+            messages = build_text_prompt_fn(
+                response_text=response_text,
+                categories_str=chunk_categories_str,
+                survey_question_context=survey_question_context,
+                examples_text=examples_text,
+                chain_of_thought=chain_of_thought,
+                context_prompt=context_prompt,
+                step_back_prompt=step_back_prompt,
+                stepback_insights=stepback_insights,
+                model_name=cfg["model"],
+                multi_label=multi_label,
+            )
+            reply, error = client.complete(
+                messages=messages,
+                json_schema=chunk_json_schema,
+                creativity=effective_creativity,
+                thinking_budget=thinking_budget if cfg["provider"] in thinking_providers else None,
+                max_retries=max_retries,
+            )
+            if error:
+                return ('{"1":"e"}', error)
+
+            json_result = extract_json(reply)
+            json_result = formatter_fallback_fn(json_result, reply, chunk_cats)
+
+            # Run Chain of Verification if enabled
+            if chain_of_verification:
+                step2, step3, step4 = build_cove_prompts_fn(
+                    chunk_cove_task, response_text
+                )
+                json_result = run_cove_fn(
+                    client=client,
+                    initial_reply=json_result,
+                    step2_prompt=step2,
+                    step3_prompt=step3,
+                    step4_prompt=step4,
+                    json_schema=chunk_json_schema,
+                    creativity=effective_creativity,
+                    max_retries=max_retries,
+                )
+                json_result = formatter_fallback_fn(json_result, json_result, chunk_cats)
+
+            return (json_result, None)

cat_stack/_embeddings.py

@@ -0,0 +1,189 @@
+"""
+Embedding-based similarity scores for CatLLM.
+
+Uses a local sentence-transformer model (BAAI/bge-small-en-v1.5, 33M params,
+~130MB) to compute cosine similarity between each input text and each category.
+Scores are independent per (text, category) pair — no softmax across categories,
+since this is multi-label classification.
+
+The embeddings feature is opt-in via embeddings=True on classify(). It adds
+`_similarity` columns alongside the existing binary 0/1 classification columns.
+
+Requires: pip install cat-llm[embeddings]
+"""
+
+import pandas as pd
+
+_EMBEDDING_MODEL_NAME = "BAAI/bge-small-en-v1.5"
+
+
+def _check_dependencies():
+    """Check that sentence-transformers is installed."""
+    try:
+        import sentence_transformers  # noqa: F401
+    except ImportError:
+        raise ImportError(
+            "The embeddings feature requires sentence-transformers.\n"
+            "Install with: pip install cat-llm[embeddings]\n"
+            "  (requires: sentence-transformers, which pulls in torch and transformers)"
+        )
+
+
+def _is_model_cached() -> bool:
+    """Check if the embedding model is already in the HuggingFace cache."""
+    try:
+        from huggingface_hub import try_to_load_from_cache
+        result = try_to_load_from_cache(_EMBEDDING_MODEL_NAME, "config.json")
+        return result is not None and not isinstance(result, type(None))
+    except Exception:
+        return False
+
+
+def ensure_embeddings_available() -> bool:
+    """
+    Ensure the embedding model is available, prompting to download if needed.
+
+    Returns:
+        True if the model is ready to use, False if user declined download.
+    """
+    _check_dependencies()
+
+    if _is_model_cached():
+        return True
+
+    print(
+        "\n[CatLLM] The embedding model (~130MB) will be downloaded from\n"
+        f"  HuggingFace Hub ({_EMBEDDING_MODEL_NAME}).\n"
+        "  This is a one-time download — the model is cached locally after."
+    )
+    try:
+        answer = input("  Continue? (Y/n): ").strip().lower()
+    except (EOFError, KeyboardInterrupt):
+        answer = "n"
+
+    if answer in ("", "y", "yes"):
+        return True
+    else:
+        print("  -> Embedding scores disabled for this run.\n")
+        return False
+
+
+def load_embedding_model():
+    """
+    Load and return the sentence-transformer embedding model.
+
+    Returns:
+        SentenceTransformer model instance.
+    """
+    _check_dependencies()
+
+    from sentence_transformers import SentenceTransformer
+
+    print(f"[CatLLM] Loading embedding model ({_EMBEDDING_MODEL_NAME})...")
+    model = SentenceTransformer(_EMBEDDING_MODEL_NAME)
+    print("[CatLLM] Embedding model ready.")
+    return model
+
+
+def compute_embedding_scores(texts, categories, model, category_descriptions=None):
+    """
+    Compute cosine similarity scores between texts and categories.
+
+    Each (text, category) score is independent — no softmax across categories.
+    Raw cosine similarity is rescaled from [-1, 1] to [0, 1] via (sim + 1) / 2.
+
+    Args:
+        texts: List of input text strings.
+        categories: List of category name strings.
+        model: Loaded SentenceTransformer model.
+        category_descriptions: Optional dict mapping category names to richer
+            descriptions for embedding (e.g., {"Past_Support": "References to
+            help received from family in the past"}).
+
+    Returns:
+        Dict mapping "category_N_similarity" -> list of float scores, where N
+        is 1-indexed to match the existing classification column naming.
+    """
+    from sentence_transformers import util
+
+    # Convert NaN/None to empty string
+    clean_texts = [str(t) if pd.notna(t) else "" for t in texts]
+
+    # Build category strings for embedding
+    cat_strings = []
+    for cat in categories:
+        if category_descriptions and cat in category_descriptions:
+            cat_strings.append(f"{cat}: {category_descriptions[cat]}")
+        else:
+            cat_strings.append(cat)
+
+    # Encode all texts and categories
+    text_embeddings = model.encode(clean_texts, normalize_embeddings=True,
+                                   show_progress_bar=len(clean_texts) > 100)
+    cat_embeddings = model.encode(cat_strings, normalize_embeddings=True)
+
+    # Compute cosine similarity matrix: (num_texts, num_categories)
+    sim_matrix = util.cos_sim(text_embeddings, cat_embeddings)
+
+    # Rescale from [-1, 1] to [0, 1]
+    scores = (sim_matrix + 1) / 2
+
+    # Build output dict
+    result = {}
+    for i, _cat in enumerate(categories):
+        col_name = f"category_{i + 1}_similarity"
+        result[col_name] = [round(float(scores[row][i]), 4) for row in range(len(clean_texts))]
+
+    return result
+
+
+def apply_embedding_scores(df, categories, embedding_model, category_descriptions=None):
+    """
+    Insert embedding similarity columns into a result DataFrame.
+
+    For each category N, a `category_N_similarity` column is inserted after the
+    last existing column that belongs to that category number.
+
+    Args:
+        df: Result DataFrame from classify (single-model or ensemble).
+        categories: List of category name strings.
+        embedding_model: Loaded SentenceTransformer model.
+        category_descriptions: Optional dict mapping category names to descriptions.
+
+    Returns:
+        DataFrame with `_similarity` columns inserted.
+    """
+    # Find the text column to use for embedding
+    if "input_data" in df.columns:
+        texts = df["input_data"].tolist()
+    else:
+        # Fallback: use first column
+        texts = df.iloc[:, 0].tolist()
+
+    scores = compute_embedding_scores(texts, categories, embedding_model,
+                                      category_descriptions)
+
+    # Insert each _similarity column after the last column for that category number
+    result_df = df.copy()
+    for i in range(len(categories)):
+        prob_col = f"category_{i + 1}_similarity"
+        prob_values = scores[prob_col]
+
+        # Find the last column that starts with "category_{N}_" or equals "category_{N}"
+        # Use exact match on the number to avoid category_1 matching category_10
+        cat_prefix = f"category_{i + 1}_"
+        cat_exact = f"category_{i + 1}"
+
+        last_pos = -1
+        for col_idx, col_name in enumerate(result_df.columns):
+            if col_name == cat_exact or col_name.startswith(cat_prefix):
+                last_pos = col_idx
+
+        if last_pos >= 0:
+            # Insert after the last matching column
+            result_df.insert(last_pos + 1, prob_col, prob_values)
+        else:
+            # No matching column found — append at the end
+            result_df[prob_col] = prob_values
+
+    return result_df