cat-stack 0.1.0__py3-none-any.whl

cat_stack/image_functions.py

@@ -0,0 +1,2078 @@
+import warnings
+
+from .text_functions import _detect_model_source
+from .calls.image_stepback import get_image_stepback_insight
+
+# Exported names (excludes deprecated image_multi_class)
+__all__ = [
+    "_load_image_files",
+    "_encode_image",
+    "image_score_drawing",
+    "image_features",
+    "explore_image_categories",
+]
+from .calls.image_CoVe import (
+    image_chain_of_verification_openai,
+    image_chain_of_verification_anthropic,
+    image_chain_of_verification_google,
+    image_chain_of_verification_mistral
+)
+
+
+def _load_image_files(image_input):
+    """Load image files from directory path, single file path, or return list as-is."""
+    import os
+    import glob
+
+    image_extensions = [
+        '*.png', '*.jpg', '*.jpeg',
+        '*.gif', '*.webp', '*.svg', '*.svgz', '*.avif', '*.apng',
+        '*.tif', '*.tiff', '*.bmp',
+        '*.heif', '*.heic', '*.ico',
+        '*.psd'
+    ]
+
+    if isinstance(image_input, list):
+        image_files = image_input
+        print(f"Provided a list of {len(image_input)} images.")
+    elif os.path.isfile(image_input):
+        # Single file path
+        image_files = [image_input]
+        print(f"Provided 1 image file.")
+    elif os.path.isdir(image_input):
+        # Directory path - glob for images
+        image_files = []
+        for ext in image_extensions:
+            image_files.extend(glob.glob(os.path.join(image_input, ext)))
+        print(f"Found {len(image_files)} images in directory.")
+    else:
+        raise FileNotFoundError(f"Image input not found: {image_input}")
+
+    return image_files
+
+
+def _encode_image(img_path):
+    """Encode an image file to base64. Returns (encoded_data, extension, is_valid)."""
+    import os
+    import base64
+    from pathlib import Path
+
+    if img_path is None or not os.path.exists(img_path):
+        return None, None, False
+
+    if os.path.isdir(img_path):
+        return None, None, False
+
+    try:
+        with open(img_path, "rb") as f:
+            encoded = base64.b64encode(f.read()).decode("utf-8")
+        ext = Path(img_path).suffix.lstrip(".").lower()
+        if ext == "jpg":
+            ext = "jpeg"
+        return encoded, ext, True
+    except Exception as e:
+        print(f"Error encoding image: {e}")
+        return None, None, False
+
+
+# image multi-class (binary) function
+def image_multi_class(
+    image_description,
+    image_input,
+    categories,
+    api_key,
+    user_model="gpt-4o",
+    creativity=None,
+    safety=False,
+    chain_of_verification=False,
+    chain_of_thought=True,
+    step_back_prompt=False,
+    context_prompt=False,
+    thinking_budget=0,
+    example1=None,
+    example2=None,
+    example3=None,
+    example4=None,
+    example5=None,
+    example6=None,
+    filename=None,
+    save_directory=None,
+    model_source="auto"
+):
+    """
+    Classify images using LLMs.
+
+    .. deprecated::
+        Use :func:`cat_stack.classify` instead. This function will be removed in a future version.
+    """
+    warnings.warn(
+        "image_multi_class() is deprecated and will be removed in a future version. "
+        "Use cat_stack.classify() instead, which auto-detects image input.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+
+    import os
+    import json
+    import pandas as pd
+    import regex
+    import time
+    from tqdm import tqdm
+
+    if save_directory is not None and not os.path.isdir(save_directory):
+        raise FileNotFoundError(f"Directory {save_directory} doesn't exist")
+
+    model_source = _detect_model_source(user_model, model_source)
+
+    image_files = _load_image_files(image_input)
+
+    # Handle "auto" categories - extract categories first
+    if categories == "auto":
+        if not image_description:
+            raise ValueError("image_description is required when using categories='auto'")
+
+        print("\nAuto-extracting categories from images...")
+        auto_result = explore_image_categories(
+            image_input=image_input,
+            api_key=api_key,
+            image_description=image_description,
+            user_model=user_model,
+            model_source=model_source,
+            creativity=creativity
+        )
+        categories = auto_result["top_categories"]
+        print(f"Extracted {len(categories)} categories: {categories}\n")
+
+    categories_str = "\n".join(f"{i + 1}. {cat}" for i, cat in enumerate(categories))
+    cat_num = len(categories)
+    category_dict = {str(i+1): "0" for i in range(cat_num)}
+    example_JSON = json.dumps(category_dict, indent=4)
+
+    print(f"\nCategories to classify by {model_source} {user_model}:")
+    for i, cat in enumerate(categories, 1):
+        print(f"{i}. {cat}")
+
+    # Build examples text from provided examples
+    examples = [example1, example2, example3, example4, example5, example6]
+    examples = [ex for ex in examples if ex is not None]
+    if examples:
+        examples_text = "Here are some examples of how to categorize:\n" + "\n".join(examples)
+    else:
+        examples_text = ""
+
+    # Helper function for CoVe
+    def remove_numbering(line):
+        line = line.strip()
+        if line.startswith('- '):
+            return line[2:].strip()
+        if line.startswith('• '):
+            return line[2:].strip()
+        if line and line[0].isdigit():
+            i = 0
+            while i < len(line) and line[i].isdigit():
+                i += 1
+            if i < len(line) and line[i] in '.':
+                return line[i+1:].strip()
+            elif i < len(line) and line[i] in ')':
+                return line[i+1:].strip()
+        return line
+
+    # Step-back insight initialization
+    if step_back_prompt:
+        stepback = f"""What are the key visual features or patterns that typically indicate the presence of these categories in images showing "{image_description}"?
+
+Categories to consider:
+{categories_str}
+
+Provide a brief analysis of what visual cues to look for when categorizing such images."""
+
+        stepback_insight, step_back_added = get_image_stepback_insight(
+            model_source, stepback, api_key, user_model, creativity
+        )
+    else:
+        stepback_insight = None
+        step_back_added = False
+
+    link1 = []
+    extracted_jsons = []
+
+    def _build_base_prompt_text():
+        """Build the base text portion of the prompt."""
+        if chain_of_thought:
+            base_text = (
+                f"You are an image-tagging assistant.\n"
+                f"Task ► Examine the attached image and decide, **for each category below**, "
+                f"whether it is PRESENT (1) or NOT PRESENT (0).\n\n"
+                f"Image is expected to show: {image_description}\n\n"
+                f"Categories:\n{categories_str}\n\n"
+                f"Let's analyze step by step:\n"
+                f"1. First, identify the key visual elements in the image\n"
+                f"2. Then, match each element to the relevant categories\n"
+                f"3. Finally, assign 1 to matching categories and 0 to non-matching categories\n\n"
+                f"{examples_text}\n\n"
+                f"Output format ► Respond with **only** a JSON object whose keys are the "
+                f"quoted category numbers ('1', '2', …) and whose values are 1 or 0. "
+                f"No additional keys, comments, or text.\n\n"
+                f"Example (three categories):\n"
+                f"{example_JSON}"
+            )
+        else:
+            base_text = (
+                f"You are an image-tagging assistant.\n"
+                f"Task ► Examine the attached image and decide, **for each category below**, "
+                f"whether it is PRESENT (1) or NOT PRESENT (0).\n\n"
+                f"Image is expected to show: {image_description}\n\n"
+                f"Categories:\n{categories_str}\n\n"
+                f"{examples_text}\n\n"
+                f"Output format ► Respond with **only** a JSON object whose keys are the "
+                f"quoted category numbers ('1', '2', …) and whose values are 1 or 0. "
+                f"No additional keys, comments, or text.\n\n"
+                f"Example (three categories):\n"
+                f"{example_JSON}"
+            )
+
+        if context_prompt:
+            context = (
+                "You are an expert visual analyst specializing in image categorization. "
+                "Apply multi-label classification based on explicit and implicit visual cues. "
+                "When uncertain, prioritize precision over recall.\n\n"
+            )
+            base_text = context + base_text
+
+        return base_text
+
+    def _build_cove_prompts(base_prompt_text):
+        """Build chain of verification prompts for images."""
+        step2_prompt = f"""You provided this initial categorization:
+<<INITIAL_REPLY>>
+
+Original task: {base_prompt_text}
+
+Generate a focused list of 3-5 verification questions to fact-check your categorization. Each question should:
+- Be concise and specific (one sentence)
+- Address a distinct visual element or category assignment
+- Be answerable by re-examining the image
+
+Focus on verifying:
+- Whether each category assignment matches what's visible in the image
+- Whether any visual elements were missed or misinterpreted
+- Whether there are any logical inconsistencies
+
+Provide only the verification questions as a numbered list."""
+
+        step3_prompt = f"""Re-examine the attached image and answer the following verification question.
+
+Image description: {image_description}
+
+Verification question: <<QUESTION>>
+
+Provide a brief, direct answer (1-2 sentences maximum) based on what you observe in the image.
+
+Answer:"""
+
+        step4_prompt = f"""Original task: {base_prompt_text}
+Initial categorization:
+<<INITIAL_REPLY>>
+Verification questions and answers:
+<<VERIFICATION_QA>>
+Based on this verification, provide the final corrected categorization.
+If no categories are present, assign "0" to all categories.
+Provide the final categorization in the same JSON format:"""
+
+        return step2_prompt, step3_prompt, step4_prompt
+
+    def _build_prompt_openai_mistral(encoded, ext, base_text):
+        """Build prompt for OpenAI/Mistral format."""
+        encoded_image = f"data:image/{ext};base64,{encoded}"
+        return [
+            {"type": "text", "text": base_text},
+            {"type": "image_url", "image_url": {"url": encoded_image, "detail": "high"}},
+        ]
+
+    def _build_prompt_anthropic(encoded, ext, base_text):
+        """Build prompt for Anthropic format."""
+        media_type = f"image/{ext}" if ext else "image/jpeg"
+        return [
+            {"type": "text", "text": base_text},
+            {
+                "type": "image",
+                "source": {
+                    "type": "base64",
+                    "media_type": media_type,
+                    "data": encoded
+                }
+            }
+        ]
+
+    def _build_prompt_google(encoded, ext, base_text):
+        """Build prompt for Google format."""
+        return {
+            "text_prompt": base_text,
+            "image_data": encoded,
+            "mime_type": f"image/{ext}" if ext else "image/jpeg"
+        }
+
+    def _call_openai_compatible(prompt, step2_prompt, step3_prompt, step4_prompt, image_content):
+        """Handle OpenAI-compatible API calls (OpenAI, Perplexity, HuggingFace, xAI)."""
+        import requests as req
+
+        # Determine the base URL based on model source
+        if model_source == "huggingface":
+            from cat_stack.text_functions import _detect_huggingface_endpoint
+            base_url = _detect_huggingface_endpoint(api_key, user_model)
+        elif model_source == "huggingface-together":
+            base_url = "https://router.huggingface.co/together/v1"
+        elif model_source == "perplexity":
+            base_url = "https://api.perplexity.ai"
+        elif model_source == "xai":
+            base_url = "https://api.x.ai/v1"
+        else:
+            base_url = "https://api.openai.com/v1"
+
+        endpoint = f"{base_url}/chat/completions"
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {api_key}"
+        }
+
+        max_retries = 8
+        delay = 2
+
+        for attempt in range(max_retries):
+            try:
+                # Build messages with optional stepback
+                messages = []
+                if step_back_prompt and step_back_added:
+                    messages.append({'role': 'user', 'content': stepback})
+                    messages.append({'role': 'assistant', 'content': stepback_insight})
+                messages.append({'role': 'user', 'content': prompt})
+
+                payload = {
+                    "model": user_model,
+                    "messages": messages,
+                }
+                if creativity is not None:
+                    payload["temperature"] = creativity
+
+                response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                response.raise_for_status()
+                result = response.json()
+                reply = result["choices"][0]["message"]["content"]
+
+                if chain_of_verification:
+                    reply = image_chain_of_verification_openai(
+                        initial_reply=reply,
+                        step2_prompt=step2_prompt,
+                        step3_prompt=step3_prompt,
+                        step4_prompt=step4_prompt,
+                        client=None,  # Not used anymore, CoVe needs refactoring too
+                        user_model=user_model,
+                        creativity=creativity,
+                        remove_numbering=remove_numbering,
+                        image_content=image_content
+                    )
+
+                return reply, None
+
+            except req.exceptions.HTTPError as e:
+                error_str = str(e).lower()
+                status_code = e.response.status_code if e.response else None
+
+                if status_code == 400 and "json_validate_failed" in error_str and attempt < max_retries - 1:
+                    wait_time = delay * (2 ** attempt)
+                    print(f"⚠️ JSON validation failed. Attempt {attempt + 1}/{max_retries}")
+                    print(f"Retrying in {wait_time}s...")
+                    time.sleep(wait_time)
+                elif status_code == 404:
+                    raise ValueError(f"❌ Model '{user_model}' on {model_source} not found. Please check the model name and try again.") from e
+                elif status_code in [500, 502, 503, 504] and attempt < max_retries - 1:
+                    wait_time = delay * (2 ** attempt)
+                    print(f"Attempt {attempt + 1} failed with error: {e}")
+                    print(f"Retrying in {wait_time}s...")
+                    time.sleep(wait_time)
+                else:
+                    print(f"❌ Failed after {max_retries} attempts: {e}")
+                    return """{"1":"e"}""", f"Error processing input: {e}"
+
+            except Exception as e:
+                if ("500" in str(e) or "504" in str(e)) and attempt < max_retries - 1:
+                    wait_time = delay * (2 ** attempt)
+                    print(f"Attempt {attempt + 1} failed with error: {e}")
+                    print(f"Retrying in {wait_time}s...")
+                    time.sleep(wait_time)
+                else:
+                    print(f"❌ Failed after {max_retries} attempts: {e}")
+                    return """{"1":"e"}""", f"Error processing input: {e}"
+
+        return """{"1":"e"}""", "Max retries exceeded"
+
+    def _call_anthropic(prompt, step2_prompt, step3_prompt, step4_prompt, image_content):
+        """Handle Anthropic API calls using direct HTTP requests."""
+        import requests as req
+
+        endpoint = "https://api.anthropic.com/v1/messages"
+        headers = {
+            "Content-Type": "application/json",
+            "x-api-key": api_key,
+            "anthropic-version": "2023-06-01"
+        }
+
+        try:
+            # Build messages with optional stepback
+            messages = []
+            if step_back_prompt and step_back_added:
+                messages.append({'role': 'user', 'content': stepback})
+                messages.append({'role': 'assistant', 'content': stepback_insight})
+            messages.append({'role': 'user', 'content': prompt})
+
+            payload = {
+                "model": user_model,
+                "max_tokens": 1024,
+                "messages": messages,
+            }
+            if creativity is not None:
+                payload["temperature"] = creativity
+
+            response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+            response.raise_for_status()
+            result = response.json()
+
+            content = result.get("content", [])
+            if content and content[0].get("type") == "text":
+                reply = content[0].get("text", "")
+            else:
+                return """{"1":"e"}""", "No text content in response"
+
+            if chain_of_verification:
+                reply = image_chain_of_verification_anthropic(
+                    initial_reply=reply,
+                    step2_prompt=step2_prompt,
+                    step3_prompt=step3_prompt,
+                    step4_prompt=step4_prompt,
+                    client=None,  # No longer using SDK client
+                    user_model=user_model,
+                    creativity=creativity,
+                    remove_numbering=remove_numbering,
+                    image_content=image_content,
+                    api_key=api_key  # Pass api_key for HTTP calls
+                )
+
+            return reply, None
+
+        except req.exceptions.HTTPError as e:
+            if e.response is not None and e.response.status_code == 404:
+                raise ValueError(f"❌ Model '{user_model}' on {model_source} not found. Please check the model name and try again.") from e
+            print(f"An error occurred: {e}")
+            return """{"1":"e"}""", f"Error processing input: {e}"
+        except Exception as e:
+            print(f"An error occurred: {e}")
+            return """{"1":"e"}""", f"Error processing input: {e}"
+
+    def _call_google(prompt_data, step2_prompt, step3_prompt, step4_prompt, base_prompt_text):
+        """Handle Google API calls."""
+        import requests
+
+        def make_google_request(url, headers, payload, max_retries=8):
+            for attempt in range(max_retries):
+                try:
+                    response = requests.post(url, headers=headers, json=payload)
+                    response.raise_for_status()
+                    return response.json()
+                except requests.exceptions.HTTPError as e:
+                    status_code = e.response.status_code
+                    retryable_errors = [429, 500, 502, 503, 504]
+
+                    if status_code in retryable_errors and attempt < max_retries - 1:
+                        wait_time = 10 * (2 ** attempt) if status_code == 429 else 2 * (2 ** attempt)
+                        error_type = "Rate limited" if status_code == 429 else f"Server error {status_code}"
+                        print(f"⚠️ {error_type}. Attempt {attempt + 1}/{max_retries}")
+                        print(f"Retrying in {wait_time}s...")
+                        time.sleep(wait_time)
+                    else:
+                        raise
+
+        url = f"https://generativelanguage.googleapis.com/v1beta/models/{user_model}:generateContent"
+        headers = {
+            "x-goog-api-key": api_key,
+            "Content-Type": "application/json"
+        }
+
+        # Build parts with optional stepback context
+        parts = []
+        if step_back_prompt and step_back_added:
+            parts.append({"text": f"Context from step-back analysis:\n{stepback_insight}\n\n"})
+        parts.append({"text": prompt_data["text_prompt"]})
+        parts.append({
+            "inline_data": {
+                "mime_type": prompt_data["mime_type"],
+                "data": prompt_data["image_data"]
+            }
+        })
+
+        payload = {
+            "contents": [{"parts": parts}],
+            "generationConfig": {
+                "responseMimeType": "application/json",
+                **({"temperature": creativity} if creativity is not None else {}),
+                **({"thinkingConfig": {"thinkingBudget": thinking_budget}} if thinking_budget else {})
+            }
+        }
+
+        try:
+            result = make_google_request(url, headers, payload)
+
+            if "candidates" in result and result["candidates"]:
+                reply = result["candidates"][0]["content"]["parts"][0]["text"]
+            else:
+                return "No response generated", None
+
+            if chain_of_verification:
+                reply = image_chain_of_verification_google(
+                    initial_reply=reply,
+                    prompt=base_prompt_text,
+                    step2_prompt=step2_prompt,
+                    step3_prompt=step3_prompt,
+                    step4_prompt=step4_prompt,
+                    url=url,
+                    headers=headers,
+                    creativity=creativity,
+                    remove_numbering=remove_numbering,
+                    make_google_request=make_google_request,
+                    image_data=prompt_data["image_data"],
+                    mime_type=prompt_data["mime_type"]
+                )
+
+            return reply, None
+
+        except requests.exceptions.HTTPError as e:
+            if e.response.status_code == 404:
+                raise ValueError(f"❌ Model '{user_model}' not found. Please check the model name and try again.") from e
+            elif e.response.status_code in [401, 403]:
+                raise ValueError(f"❌ Authentication failed. Please check your Google API key.") from e
+            else:
+                print(f"HTTP error occurred: {e}")
+                return """{"1":"e"}""", f"Error processing input: {e}"
+        except Exception as e:
+            print(f"An error occurred: {e}")
+            return """{"1":"e"}""", f"Error processing input: {e}"
+
+    def _call_mistral(prompt, step2_prompt, step3_prompt, step4_prompt, image_content):
+        """Handle Mistral API calls - uses requests directly."""
+        import requests as req
+
+        endpoint = "https://api.mistral.ai/v1/chat/completions"
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {api_key}"
+        }
+
+        max_retries = 8
+        delay = 2
+
+        for attempt in range(max_retries):
+            try:
+                # Build messages with optional stepback
+                messages = []
+                if step_back_prompt and step_back_added:
+                    messages.append({'role': 'user', 'content': stepback})
+                    messages.append({'role': 'assistant', 'content': stepback_insight})
+                messages.append({'role': 'user', 'content': prompt})
+
+                payload = {
+                    "model": user_model,
+                    "messages": messages,
+                }
+                if creativity is not None:
+                    payload["temperature"] = creativity
+
+                response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                response.raise_for_status()
+                result = response.json()
+                reply = result["choices"][0]["message"]["content"]
+
+                if chain_of_verification:
+                    reply = image_chain_of_verification_mistral(
+                        initial_reply=reply,
+                        step2_prompt=step2_prompt,
+                        step3_prompt=step3_prompt,
+                        step4_prompt=step4_prompt,
+                        client=None,  # Not used anymore, CoVe needs refactoring too
+                        user_model=user_model,
+                        creativity=creativity,
+                        remove_numbering=remove_numbering,
+                        image_content=image_content
+                    )
+
+                return reply, None
+
+            except req.exceptions.HTTPError as e:
+                error_str = str(e).lower()
+                status_code = e.response.status_code if e.response else None
+
+                if status_code == 404 or "invalid_model" in error_str or "invalid model" in error_str:
+                    raise ValueError(f"❌ Model '{user_model}' not found.") from e
+                elif status_code == 401 or "unauthorized" in error_str:
+                    raise ValueError(f"❌ Authentication failed. Please check your Mistral API key.") from e
+                elif status_code in [500, 502, 503, 504] and attempt < max_retries - 1:
+                    wait_time = delay * (2 ** attempt)
+                    print(f"⚠️ Server error {status_code}. Attempt {attempt + 1}/{max_retries}")
+                    print(f"Retrying in {wait_time}s...")
+                    time.sleep(wait_time)
+                else:
+                    print(f"❌ Failed after {max_retries} attempts: {e}")
+                    return """{"1":"e"}""", f"Error processing input: {e}"
+
+            except Exception as e:
+                print(f"❌ Unexpected error: {e}")
+                return """{"1":"e"}""", f"Error processing input: {e}"
+
+        return """{"1":"e"}""", "Max retries exceeded"
+
+    def _process_single_image(img_path):
+        """Process a single image and return (reply, error_msg)."""
+        encoded, ext, is_valid = _encode_image(img_path)
+
+        if not is_valid:
+            return None, "Invalid image path or encoding failed"
+
+        base_prompt_text = _build_base_prompt_text()
+
+        if chain_of_verification:
+            step2_prompt, step3_prompt, step4_prompt = _build_cove_prompts(base_prompt_text)
+        else:
+            step2_prompt = step3_prompt = step4_prompt = None
+
+        if model_source in ["openai", "perplexity", "huggingface", "xai"]:
+            prompt = _build_prompt_openai_mistral(encoded, ext, base_prompt_text)
+            # Image content for CoVe (just the image part)
+            encoded_image = f"data:image/{ext};base64,{encoded}"
+            image_content = {"type": "image_url", "image_url": {"url": encoded_image, "detail": "high"}}
+            return _call_openai_compatible(prompt, step2_prompt, step3_prompt, step4_prompt, image_content)
+
+        elif model_source == "anthropic":
+            prompt = _build_prompt_anthropic(encoded, ext, base_prompt_text)
+            media_type = f"image/{ext}" if ext else "image/jpeg"
+            image_content = {
+                "type": "image",
+                "source": {"type": "base64", "media_type": media_type, "data": encoded}
+            }
+            return _call_anthropic(prompt, step2_prompt, step3_prompt, step4_prompt, image_content)
+
+        elif model_source == "google":
+            prompt_data = _build_prompt_google(encoded, ext, base_prompt_text)
+            return _call_google(prompt_data, step2_prompt, step3_prompt, step4_prompt, base_prompt_text)
+
+        elif model_source == "mistral":
+            prompt = _build_prompt_openai_mistral(encoded, ext, base_prompt_text)
+            encoded_image = f"data:image/{ext};base64,{encoded}"
+            image_content = {"type": "image_url", "image_url": {"url": encoded_image, "detail": "high"}}
+            return _call_mistral(prompt, step2_prompt, step3_prompt, step4_prompt, image_content)
+
+        else:
+            raise ValueError("Unknown source! Choose from OpenAI, Anthropic, Perplexity, Google, xAI, Huggingface, or Mistral")
+
+    def _extract_json(reply):
+        """Extract JSON from model reply."""
+        if reply is None:
+            return """{"1":"e"}"""
+
+        if reply == "invalid image path":
+            return """{"no_valid_path": 1}"""
+
+        extracted_json = regex.findall(r'\{(?:[^{}]|(?R))*\}', reply, regex.DOTALL)
+        if extracted_json:
+            return extracted_json[0].replace('[', '').replace(']', '').replace('\n', '').replace(" ", '').replace("  ", '')
+        else:
+            print("""{"1":"e"}""")
+            return """{"1":"e"}"""
+
+    # Main processing loop
+    for idx, img_path in enumerate(tqdm(image_files, desc="Categorizing images")):
+        if img_path is None:
+            link1.append("Skipped NaN input")
+            extracted_jsons.append("""{"no_valid_image": 1}""")
+            continue
+
+        reply, error_msg = _process_single_image(img_path)
+
+        if error_msg:
+            link1.append(error_msg)
+            if "Invalid image" in error_msg:
+                extracted_jsons.append("""{"no_valid_path": 1}""")
+            else:
+                extracted_jsons.append(_extract_json(reply))
+        else:
+            link1.append(reply)
+            extracted_jsons.append(_extract_json(reply))
+
+        # --- Safety Save ---
+        if safety:
+            if filename is None:
+                raise TypeError("filename is required when using safety. Please provide the filename.")
+
+            normalized_data_list = []
+            for json_str in extracted_jsons:
+                try:
+                    parsed_obj = json.loads(json_str)
+                    normalized_data_list.append(pd.json_normalize(parsed_obj))
+                except json.JSONDecodeError:
+                    normalized_data_list.append(pd.DataFrame({"1": ["e"]}))
+            normalized_data = pd.concat(normalized_data_list, ignore_index=True)
+
+            temp_df = pd.DataFrame({
+                'image_input': image_files[:idx+1],
+                'model_response': link1,
+                'json': extracted_jsons
+            })
+            temp_df = pd.concat([temp_df, normalized_data], axis=1)
+
+            save_path = os.path.join(save_directory, filename) if save_directory else filename
+            temp_df.to_csv(save_path, index=False)
+
+    # --- Final DataFrame ---
+    normalized_data_list = []
+    for json_str in extracted_jsons:
+        try:
+            parsed_obj = json.loads(json_str)
+            normalized_data_list.append(pd.json_normalize(parsed_obj))
+        except json.JSONDecodeError:
+            normalized_data_list.append(pd.DataFrame({"1": ["e"]}))
+    normalized_data = pd.concat(normalized_data_list, ignore_index=True)
+
+    categorized_data = pd.DataFrame({
+        'image_input': pd.Series(image_files),
+        'model_response': pd.Series(link1).reset_index(drop=True),
+        'json': pd.Series(extracted_jsons).reset_index(drop=True)
+    })
+    categorized_data = pd.concat([categorized_data, normalized_data], axis=1)
+    categorized_data = categorized_data.rename(columns=lambda x: f'category_{x}' if str(x).isdigit() else x)
+
+    # Identify rows with invalid strings (like "e")
+    cat_cols = [col for col in categorized_data.columns if col.startswith('category_')]
+    has_invalid_strings = categorized_data[cat_cols].apply(
+        lambda col: pd.to_numeric(col, errors='coerce').isna() & col.notna()
+    ).any(axis=1)
+
+    categorized_data['processing_status'] = (~has_invalid_strings).map({True: 'success', False: 'error'})
+    categorized_data.loc[has_invalid_strings, cat_cols] = pd.NA
+
+    for col in cat_cols:
+        categorized_data[col] = pd.to_numeric(categorized_data[col], errors='coerce')
+
+    categorized_data.loc[~has_invalid_strings, cat_cols] = (
+        categorized_data.loc[~has_invalid_strings, cat_cols].fillna(0)
+    )
+    categorized_data[cat_cols] = categorized_data[cat_cols].astype('Int64')
+
+    # Create categories_id (comma-separated binary values for each category)
+    categorized_data['categories_id'] = categorized_data[cat_cols].apply(
+        lambda x: ','.join(x.dropna().astype(int).astype(str)), axis=1
+    )
+
+    if filename:
+        save_path = os.path.join(save_directory, filename) if save_directory else filename
+        categorized_data.to_csv(save_path, index=False)
+
+    return categorized_data
+
+
+# image score function
+def image_score_drawing(
+    reference_image_description,
+    image_input,
+    reference_image,
+    api_key,
+    columns="numbered",
+    user_model="gpt-4o-2024-11-20",
+    creativity=None,
+    to_csv=False,
+    safety=False,
+    filename="categorized_data.csv",
+    save_directory=None,
+    model_source="OpenAI"
+):
+    import os
+    import json
+    import pandas as pd
+    import regex
+    from tqdm import tqdm
+    import glob
+    import base64
+    from pathlib import Path
+
+    if save_directory is not None and not os.path.isdir(save_directory):
+    # Directory doesn't exist - raise an exception to halt execution
+        raise FileNotFoundError(f"Directory {save_directory} doesn't exist")
+
+    image_extensions = [
+    '*.png', '*.jpg', '*.jpeg',
+    '*.gif', '*.webp', '*.svg', '*.svgz', '*.avif', '*.apng',
+    '*.tif', '*.tiff', '*.bmp',
+    '*.heif', '*.heic', '*.ico',
+    '*.psd'
+    ]
+
+    model_source = model_source.lower() # eliminating case sensitivity
+
+    if not isinstance(image_input, list):
+        # If image_input is a filepath (string)
+        image_files = []
+        for ext in image_extensions:
+            image_files.extend(glob.glob(os.path.join(image_input, ext)))
+
+        print(f"Found {len(image_files)} images.")
+    else:
+        # If image_files is already a list
+        image_files = image_input
+        print(f"Provided a list of {len(image_input)} images.")
+
+    with open(reference_image, 'rb') as f:
+        reference = base64.b64encode(f.read()).decode('utf-8')
+        reference_image = f"data:image/{reference_image.split('.')[-1]};base64,{reference}"
+
+    link1 = []
+    extracted_jsons = []
+
+    for i, img_path in enumerate(tqdm(image_files, desc="Categorising images"), start=0):
+    # Check validity first
+        if img_path is None or not os.path.exists(img_path):
+            link1.append("Skipped NaN input or invalid path")
+            extracted_jsons.append("""{"no_valid_image": 1}""")
+            continue  # Skip the rest of the loop iteration
+
+    # Only open the file if path is valid
+        if os.path.isdir(img_path):
+            encoded = "Not a Valid Image, contains file path"
+        else:
+            try:
+                with open(img_path, "rb") as f:
+                    encoded = base64.b64encode(f.read()).decode("utf-8")
+            except Exception as e:
+                    encoded = f"Error: {str(e)}"
+    # Handle extension safely
+        if encoded.startswith("Error:") or encoded == "Not a Valid Image, contains file path":
+            encoded_image = encoded
+            valid_image = False
+
+        else:
+            ext = Path(img_path).suffix.lstrip(".").lower()
+            encoded_image = f"data:image/{ext};base64,{encoded}"
+            valid_image = True
+
+    # Handle extension safely
+        ext = Path(img_path).suffix.lstrip(".").lower()
+        encoded_image = f"data:image/{ext};base64,{encoded}"
+
+        if model_source == "openai" or model_source == "mistral":
+            prompt = [
+                {
+                    "type": "text",
+                    "text": (
+                        f"You are a visual similarity assessment system.\n"
+                        f"Task ► Compare these two images:\n"
+                        f"1. REFERENCE (left): {reference_image_description}\n"
+                        f"2. INPUT (right): User-provided drawing\n\n"
+                        f"Rating criteria:\n"
+                        f"1: No meaningful similarity (fundamentally different)\n"
+                        f"2: Barely recognizable similarity (25% match)\n"
+                        f"3: Partial match (50% key features)\n"
+                        f"4: Strong alignment (75% features)\n"
+                        f"5: Near-perfect match (90%+ similarity)\n\n"
+                        f"Output format ► Return ONLY:\n"
+                        "{\n"
+                        '  "score": [1-5],\n'
+                        '  "summary": "reason you scored"\n'
+                        "}\n\n"
+                        f"Critical rules:\n"
+                        f"- Score must reflect shape, proportions, and key details\n"
+                        f"- List only concrete matching elements from reference\n"
+                        f"- No markdown or additional text"
+                    )
+                },
+                {
+                    "type": "image_url",
+                    "image_url": {"url": reference_image, "detail": "high"}
+                },
+                {
+                    "type": "image_url",
+                    "image_url": {"url": encoded_image, "detail": "high"}
+                }
+            ]
+
+        elif model_source == "anthropic":  # Changed to elif
+            prompt = [
+                {
+                    "type": "text",
+                    "text": (
+                        f"You are a visual similarity assessment system.\n"
+                        f"Task ► Compare these two images:\n"
+                        f"1. REFERENCE (left): {reference_image_description}\n"
+                        f"2. INPUT (right): User-provided drawing\n\n"
+                        f"Rating criteria:\n"
+                        f"1: No meaningful similarity (fundamentally different)\n"
+                        f"2: Barely recognizable similarity (25% match)\n"
+                        f"3: Partial match (50% key features)\n"
+                        f"4: Strong alignment (75% features)\n"
+                        f"5: Near-perfect match (90%+ similarity)\n\n"
+                        f"Output format ► Return ONLY:\n"
+                        "{\n"
+                        '  "score": [1-5],\n'
+                        '  "summary": "reason you scored"\n'
+                        "}\n\n"
+                        f"Critical rules:\n"
+                        f"- Score must reflect shape, proportions, and key details\n"
+                        f"- List only concrete matching elements from reference\n"
+                        f"- No markdown or additional text"
+                    )
+                },
+                {
+                    "type": "image",  # Added missing type
+                    "source": {
+                        "type": "base64",
+                        "media_type": "image/png",
+                        "data": reference
+                    }
+                },
+                {
+                    "type": "image",  # Added missing type
+                    "source": {
+                        "type": "base64",
+                        "media_type": "image/jpeg",
+                        "data": encoded
+                    }
+                }
+            ]
+
+
+        if model_source == "openai":
+            import requests as req
+            endpoint = "https://api.openai.com/v1/chat/completions"
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {api_key}"
+            }
+            payload = {
+                "model": user_model,
+                "messages": [{'role': 'user', 'content': prompt}],
+            }
+            if creativity is not None:
+                payload["temperature"] = creativity
+            try:
+                response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                response.raise_for_status()
+                result = response.json()
+                reply = result["choices"][0]["message"]["content"]
+                link1.append(reply)
+            except req.exceptions.HTTPError as e:
+                if e.response and e.response.status_code == 404:
+                    raise ValueError(f"Invalid OpenAI model '{user_model}': {e}")
+                else:
+                    print(f"An error occurred: {e}")
+                    link1.append(f"Error processing input: {e}")
+            except Exception as e:
+                print(f"An error occurred: {e}")
+                link1.append(f"Error processing input: {e}")
+
+        elif model_source == "anthropic":
+            import requests as req
+            endpoint = "https://api.anthropic.com/v1/messages"
+            headers = {
+                "Content-Type": "application/json",
+                "x-api-key": api_key,
+                "anthropic-version": "2023-06-01"
+            }
+            payload = {
+                "model": user_model,
+                "max_tokens": 1024,
+                "messages": [{"role": "user", "content": prompt}],
+            }
+            if creativity is not None:
+                payload["temperature"] = creativity
+            try:
+                response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                response.raise_for_status()
+                result = response.json()
+                content = result.get("content", [])
+                if content and content[0].get("type") == "text":
+                    reply = content[0].get("text", "")
+                    link1.append(reply)
+                else:
+                    link1.append("Error processing input: No text content in response")
+            except req.exceptions.HTTPError as e:
+                if e.response is not None and e.response.status_code == 404:
+                    raise ValueError(f"Invalid Anthropic model '{user_model}': {e}")
+                else:
+                    print(f"An error occurred: {e}")
+                    link1.append(f"Error processing input: {e}")
+            except Exception as e:
+                print(f"An error occurred: {e}")
+                link1.append(f"Error processing input: {e}")
+
+        elif model_source == "mistral":
+            import requests as req
+            endpoint = "https://api.mistral.ai/v1/chat/completions"
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {api_key}"
+            }
+            payload = {
+                "model": user_model,
+                "messages": [{'role': 'user', 'content': prompt}],
+            }
+            if creativity is not None:
+                payload["temperature"] = creativity
+            try:
+                response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                response.raise_for_status()
+                result = response.json()
+                reply = result["choices"][0]["message"]["content"]
+                link1.append(reply)
+            except req.exceptions.HTTPError as e:
+                if e.response and e.response.status_code == 404:
+                    raise ValueError(f"Invalid Mistral model '{user_model}': {e}")
+                else:
+                    print(f"An error occurred: {e}")
+                    link1.append(f"Error processing input: {e}")
+            except Exception as e:
+                print(f"An error occurred: {e}")
+                link1.append(f"Error processing input: {e}")
+        #if no valid image path is provided
+        elif  valid_image == False:
+            reply = "invalid image path"
+            print("Skipped NaN input or invalid path")
+            #extracted_jsons.append("""{"no_valid_path": 1}""")
+            link1.append("Error processing input: {e}")
+        else:
+            raise ValueError("Unknown source! Choose from OpenAI, Perplexity, or Mistral")
+            # in situation that no JSON is found
+        if reply is not None:
+            if reply == "invalid image path":
+                extracted_jsons.append("""{"no_valid_path": 1}""")
+            else:
+                extracted_json = regex.findall(r'\{(?:[^{}]|(?R))*\}', reply, regex.DOTALL)
+                if extracted_json:
+                    cleaned_json = extracted_json[0].replace('[', '').replace(']', '').replace('\n', '').replace(" ", '').replace("  ", '')
+                    extracted_jsons.append(cleaned_json)
+                else:
+                    error_message = """{"1":"e"}"""
+                    extracted_jsons.append(error_message)
+                    print(error_message)
+        else:
+            error_message = """{"1":"e"}"""
+            extracted_jsons.append(error_message)
+            print(error_message)
+
+        # --- Safety Save ---
+        if safety:
+            # Save progress so far
+            temp_df = pd.DataFrame({
+                'image_input': image_files[:i+1],
+                'model_response': link1,
+                'json': extracted_jsons
+            })
+            # Normalize processed jsons so far
+            normalized_data_list = []
+            for json_str in extracted_jsons:
+                try:
+                    parsed_obj = json.loads(json_str)
+                    normalized_data_list.append(pd.json_normalize(parsed_obj))
+                except json.JSONDecodeError:
+                    normalized_data_list.append(pd.DataFrame({"1": ["e"]}))
+            normalized_data = pd.concat(normalized_data_list, ignore_index=True)
+            temp_df = pd.concat([temp_df, normalized_data], axis=1)
+            # Save to CSV
+            if save_directory is None:
+                save_directory = os.getcwd()
+            temp_df.to_csv(os.path.join(save_directory, filename), index=False)
+
+    # --- Final DataFrame ---
+    normalized_data_list = []
+    for json_str in extracted_jsons:
+        try:
+            parsed_obj = json.loads(json_str)
+            normalized_data_list.append(pd.json_normalize(parsed_obj))
+        except json.JSONDecodeError:
+            normalized_data_list.append(pd.DataFrame({"1": ["e"]}))
+    normalized_data = pd.concat(normalized_data_list, ignore_index=True)
+
+    categorized_data = pd.DataFrame({
+        'image_input': (
+            image_files.reset_index(drop=True) if isinstance(image_files, (pd.DataFrame, pd.Series))
+            else pd.Series(image_files)
+        ),
+        'link1': pd.Series(link1).reset_index(drop=True),
+        'json': pd.Series(extracted_jsons).reset_index(drop=True)
+    })
+    categorized_data = pd.concat([categorized_data, normalized_data], axis=1)
+
+    if to_csv:
+        if save_directory is None:
+            save_directory = os.getcwd()
+        categorized_data.to_csv(os.path.join(save_directory, filename), index=False)
+
+    return categorized_data
+
+# image features function
+def image_features(
+    image_description,
+    image_input,
+    features_to_extract,
+    api_key,
+    user_model="gpt-4o-2024-11-20",
+    creativity=None,
+    to_csv=False,
+    safety=False,
+    filename="categorized_data.csv",
+    save_directory=None,
+    model_source="OpenAI"
+):
+    import os
+    import json
+    import pandas as pd
+    import regex
+    from tqdm import tqdm
+    import glob
+    import base64
+    from pathlib import Path
+
+    image_extensions = [
+    '*.png', '*.jpg', '*.jpeg',
+    '*.gif', '*.webp', '*.svg', '*.svgz', '*.avif', '*.apng',
+    '*.tif', '*.tiff', '*.bmp',
+    '*.heif', '*.heic', '*.ico',
+    '*.psd'
+    ]
+
+    model_source = model_source.lower() # eliminating case sensitivity
+
+    if not isinstance(image_input, list):
+        # If image_input is a filepath (string)
+        image_files = []
+        for ext in image_extensions:
+            image_files.extend(glob.glob(os.path.join(image_input, ext)))
+
+        print(f"Found {len(image_files)} images.")
+    else:
+        # If image_files is already a list
+        image_files = image_input
+        print(f"Provided a list of {len(image_input)} images.")
+
+    categories_str = "\n".join(f"{i + 1}. {cat}" for i, cat in enumerate(features_to_extract))
+    cat_num = len(features_to_extract)
+    category_dict = {str(i+1): "0" for i in range(cat_num)}
+    example_JSON = json.dumps(category_dict, indent=4)
+
+    link1 = []
+    extracted_jsons = []
+
+    for i, img_path in enumerate(tqdm(image_files, desc="Scoring images"), start=0):
+    # Check validity first
+        if img_path is None or not os.path.exists(img_path):
+            link1.append("Skipped NaN input or invalid path")
+            extracted_jsons.append("""{"no_valid_image": 1}""")
+            continue  # Skip the rest of the loop iteration
+
+    # Only open the file if path is valid
+        if os.path.isdir(img_path):
+            encoded = "Not a Valid Image, contains file path"
+        else:
+            try:
+                with open(img_path, "rb") as f:
+                    encoded = base64.b64encode(f.read()).decode("utf-8")
+            except Exception as e:
+                    encoded = f"Error: {str(e)}"
+    # Handle extension safely
+        if encoded.startswith("Error:") or encoded == "Not a Valid Image, contains file path":
+            encoded_image = encoded
+            valid_image = False
+
+        else:
+            ext = Path(img_path).suffix.lstrip(".").lower()
+            encoded_image = f"data:image/{ext};base64,{encoded}"
+            valid_image = True
+
+        if model_source == "openai" or model_source == "mistral":
+            prompt = [
+                {
+                    "type": "text",
+                    "text": (
+                        f"You are a visual question answering assistant.\n"
+                        f"Task ► Analyze the attached image and answer these specific questions:\n\n"
+                        f"Image context: {image_description}\n\n"
+                        f"Questions to answer:\n{categories_str}\n\n"
+                        f"Output format ► Return **only** a JSON object where:\n"
+                        f"- Keys are question numbers ('1', '2', ...)\n"
+                        f"- Values are concise answers (numbers, short phrases)\n\n"
+                        f"Example for 3 questions:\n"
+                        "{\n"
+                        '  "1": "4",\n'
+                        '  "2": "blue",\n'
+                        '  "3": "yes"\n'
+                        "}\n\n"
+                        f"Important rules:\n"
+                        f"1. Answer directly - no explanations\n"
+                        f"2. Use exact numerical values when possible\n"
+                        f"3. For yes/no questions, use 'yes' or 'no'\n"
+                        f"4. Never add extra keys or formatting"
+                        ),
+                        },
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": encoded_image, "detail": "high"},
+                            },
+            ]
+        elif model_source == "anthropic":
+            prompt = [
+                {
+                    "type": "text",
+                    "text": (
+                        f"You are a visual question answering assistant.\n"
+                        f"Task ► Analyze the attached image and answer these specific questions:\n\n"
+                        f"Image context: {image_description}\n\n"
+                        f"Questions to answer:\n{categories_str}\n\n"
+                        f"Output format ► Return **only** a JSON object where:\n"
+                        f"- Keys are question numbers ('1', '2', ...)\n"
+                        f"- Values are concise answers (numbers, short phrases)\n\n"
+                        f"Example for 3 questions:\n"
+                        "{\n"
+                        '  "1": "4",\n'
+                        '  "2": "blue",\n'
+                        '  "3": "yes"\n'
+                        "}\n\n"
+                        f"Important rules:\n"
+                        f"1. Answer directly - no explanations\n"
+                        f"2. Use exact numerical values when possible\n"
+                        f"3. For yes/no questions, use 'yes' or 'no'\n"
+                        f"4. Never add extra keys or formatting"
+                    )
+                },
+                {
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": "image/jpeg",
+                        "data": encoded
+                    }
+                }
+            ]
+        if model_source == "openai":
+            import requests as req
+            endpoint = "https://api.openai.com/v1/chat/completions"
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {api_key}"
+            }
+            payload = {
+                "model": user_model,
+                "messages": [{'role': 'user', 'content': prompt}],
+            }
+            if creativity is not None:
+                payload["temperature"] = creativity
+            try:
+                response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                response.raise_for_status()
+                result = response.json()
+                reply = result["choices"][0]["message"]["content"]
+                link1.append(reply)
+            except req.exceptions.HTTPError as e:
+                if e.response and e.response.status_code == 404:
+                    raise ValueError(f"Invalid OpenAI model '{user_model}': {e}")
+                else:
+                    print(f"An error occurred: {e}")
+                    link1.append(f"Error processing input: {e}")
+            except Exception as e:
+                print(f"An error occurred: {e}")
+                link1.append(f"Error processing input: {e}")
+
+        elif model_source == "perplexity":
+            import requests as req
+            endpoint = "https://api.perplexity.ai/chat/completions"
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {api_key}"
+            }
+            payload = {
+                "model": user_model,
+                "messages": [{'role': 'user', 'content': prompt}],
+            }
+            if creativity is not None:
+                payload["temperature"] = creativity
+            try:
+                response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                response.raise_for_status()
+                result = response.json()
+                reply = result["choices"][0]["message"]["content"]
+                link1.append(reply)
+            except req.exceptions.HTTPError as e:
+                if e.response and e.response.status_code == 404:
+                    raise ValueError(f"Invalid Perplexity model '{user_model}': {e}")
+                else:
+                    print(f"An error occurred: {e}")
+                    link1.append(f"Error processing input: {e}")
+            except Exception as e:
+                print(f"An error occurred: {e}")
+                link1.append(f"Error processing input: {e}")
+
+        elif model_source == "anthropic":
+            import requests as req
+            endpoint = "https://api.anthropic.com/v1/messages"
+            headers = {
+                "Content-Type": "application/json",
+                "x-api-key": api_key,
+                "anthropic-version": "2023-06-01"
+            }
+            payload = {
+                "model": user_model,
+                "max_tokens": 1024,
+                "messages": [{"role": "user", "content": prompt}],
+            }
+            if creativity is not None:
+                payload["temperature"] = creativity
+            try:
+                response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                response.raise_for_status()
+                result = response.json()
+                content = result.get("content", [])
+                if content and content[0].get("type") == "text":
+                    reply = content[0].get("text", "")
+                    link1.append(reply)
+                else:
+                    link1.append("Error processing input: No text content in response")
+            except req.exceptions.HTTPError as e:
+                if e.response is not None and e.response.status_code == 404:
+                    raise ValueError(f"Invalid Anthropic model '{user_model}': {e}")
+                else:
+                    print(f"An error occurred: {e}")
+                    link1.append(f"Error processing input: {e}")
+            except Exception as e:
+                print(f"An error occurred: {e}")
+                link1.append(f"Error processing input: {e}")
+
+        elif model_source == "mistral":
+            import requests as req
+            endpoint = "https://api.mistral.ai/v1/chat/completions"
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {api_key}"
+            }
+            payload = {
+                "model": user_model,
+                "messages": [{'role': 'user', 'content': prompt}],
+            }
+            if creativity is not None:
+                payload["temperature"] = creativity
+            try:
+                response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                response.raise_for_status()
+                result = response.json()
+                reply = result["choices"][0]["message"]["content"]
+                link1.append(reply)
+            except req.exceptions.HTTPError as e:
+                if e.response and e.response.status_code == 404:
+                    raise ValueError(f"Invalid Mistral model '{user_model}': {e}")
+                else:
+                    print(f"An error occurred: {e}")
+                    link1.append(f"Error processing input: {e}")
+            except Exception as e:
+                print(f"An error occurred: {e}")
+                link1.append(f"Error processing input: {e}")
+
+        elif  valid_image == False:
+            print("Skipped NaN input or invalid path")
+            reply = None
+            link1.append(f"Error processing input: invalid image")
+        else:
+            raise ValueError("Unknown source! Choose from OpenAI, Anthropic, Perplexity, or Mistral")
+            # in situation that no JSON is found
+        if reply is not None:
+            extracted_json = regex.findall(r'\{(?:[^{}]|(?R))*\}', reply, regex.DOTALL)
+            if extracted_json:
+                cleaned_json = extracted_json[0].replace('[', '').replace(']', '').replace('\n', '').replace(" ", '').replace("  ", '')
+                extracted_jsons.append(cleaned_json)
+                #print(cleaned_json)
+            else:
+                error_message = """{"1":"e"}"""
+                extracted_jsons.append(error_message)
+                print(error_message)
+        else:
+            error_message = """{"1":"e"}"""
+            extracted_jsons.append(error_message)
+            #print(error_message)
+
+        # --- Safety Save ---
+        if safety:
+            #print(f"Saving CSV to: {save_directory}")
+            # Save progress so far
+            temp_df = pd.DataFrame({
+                'image_input': image_files[:i+1],
+                'link1': link1,
+                'json': extracted_jsons
+            })
+            # Normalize processed jsons so far
+            normalized_data_list = []
+            for json_str in extracted_jsons:
+                try:
+                    parsed_obj = json.loads(json_str)
+                    normalized_data_list.append(pd.json_normalize(parsed_obj))
+                except json.JSONDecodeError:
+                    normalized_data_list.append(pd.DataFrame({"1": ["e"]}))
+            normalized_data = pd.concat(normalized_data_list, ignore_index=True)
+            temp_df = pd.concat([temp_df, normalized_data], axis=1)
+            # Save to CSV
+            if save_directory is None:
+                save_directory = os.getcwd()
+            temp_df.to_csv(os.path.join(save_directory, filename), index=False)
+
+    # --- Final DataFrame ---
+    normalized_data_list = []
+    for json_str in extracted_jsons:
+        try:
+            parsed_obj = json.loads(json_str)
+            normalized_data_list.append(pd.json_normalize(parsed_obj))
+        except json.JSONDecodeError:
+            normalized_data_list.append(pd.DataFrame({"1": ["e"]}))
+    normalized_data = pd.concat(normalized_data_list, ignore_index=True)
+
+    categorized_data = pd.DataFrame({
+        'image_input': (
+            image_files.reset_index(drop=True) if isinstance(image_files, (pd.DataFrame, pd.Series))
+            else pd.Series(image_files)
+        ),
+        'model_response': pd.Series(link1).reset_index(drop=True),
+        'json': pd.Series(extracted_jsons).reset_index(drop=True)
+    })
+    categorized_data = pd.concat([categorized_data, normalized_data], axis=1)
+
+    if to_csv:
+        if save_directory is None:
+            save_directory = os.getcwd()
+        categorized_data.to_csv(os.path.join(save_directory, filename), index=False)
+
+    return categorized_data
+
+
+def explore_image_categories(
+    image_input,
+    api_key,
+    image_description="",
+    max_categories=12,
+    categories_per_chunk=10,
+    divisions=5,
+    user_model="gpt-4o",
+    creativity=None,
+    specificity="broad",
+    research_question=None,
+    mode="image",
+    filename=None,
+    model_source="auto",
+    iterations=3,
+    random_state=None,
+    progress_callback=None,
+):
+    """
+    Explore and extract common categories from a collection of images.
+
+    Modes:
+        - "image" (default): Samples random images and sends them directly to
+          a vision model for category extraction. Best for visual categorization.
+
+        - "both": Samples random images, uses vision model to describe each
+          image's content (including any text), then extracts categories from
+          those descriptions. Best for images that contain text or mixed content.
+
+    Args:
+        image_input: Path to image file, directory of images, or list of image paths
+        api_key: API key for the model provider
+        image_description: Description of what the images contain
+        max_categories: Maximum number of final categories to return
+        categories_per_chunk: Categories to extract per chunk of images
+        divisions: Number of chunks to divide images into
+        user_model: Model to use (must support vision)
+        creativity: Temperature setting (None for default)
+        specificity: "broad" or "specific" category granularity
+        research_question: Optional research context
+        mode: "image" or "both"
+        filename: Optional CSV filename to save results
+        model_source: "auto", "openai", "anthropic", "google", "mistral"
+        iterations: Number of passes over the data
+        random_state: Random seed for reproducibility
+        progress_callback: Optional callback function for progress updates.
+            Called as progress_callback(current_step, total_steps, step_label).
+
+    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
+    """
+    import os
+    import re
+    import pandas as pd
+    import numpy as np
+    from tqdm import tqdm
+
+    model_source = _detect_model_source(user_model, model_source)
+
+    # Load all images
+    image_files = _load_image_files(image_input)
+    if not image_files:
+        raise ValueError("No image files found in the specified input.")
+
+    n = len(image_files)
+    if n == 0:
+        raise ValueError("No images found.")
+
+    # Auto-adjust divisions for small datasets
+    # Images can have multiple categories each, so we can use fewer divisions
+    original_divisions = divisions
+    divisions = min(divisions, max(1, n // 2))  # At least 2 images per chunk
+    if divisions != original_divisions:
+        print(f"Auto-adjusted divisions from {original_divisions} to {divisions} for {n} images.")
+
+    # Chunk sizing - images often contain multiple categories each
+    chunk_size = int(round(max(1, n / divisions), 0))
+    # Don't reduce categories_per_chunk as aggressively for images since each image can yield many categories
+    if chunk_size < 2:
+        # Only reduce if we have very few images
+        old_categories_per_chunk = categories_per_chunk
+        categories_per_chunk = max(5, chunk_size * 4)
+        print(f"Auto-adjusted categories_per_chunk from {old_categories_per_chunk} to {categories_per_chunk} for chunk size {chunk_size}.")
+
+    print(
+        f"Exploring categories in images: '{image_description}'.\n"
+        f"          {n} total images, {categories_per_chunk * divisions} categories to extract, "
+        f"{max_categories} final categories. Mode: {mode}\n"
+    )
+
+    # RNG for reproducible sampling
+    rng = np.random.default_rng(random_state)
+
+    # Validate model_source (clients initialized per-call using requests)
+    import requests as req
+    if model_source not in ["openai", "huggingface", "huggingface-together", "xai", "anthropic", "google", "mistral"]:
+        raise ValueError(f"Unsupported model_source: {model_source}")
+
+    # Determine base URL for OpenAI-compatible providers
+    if model_source == "huggingface":
+        from cat_stack.text_functions import _detect_huggingface_endpoint
+        openai_base_url = _detect_huggingface_endpoint(api_key, user_model)
+    elif model_source == "huggingface-together":
+        openai_base_url = "https://router.huggingface.co/together/v1"
+    elif model_source == "xai":
+        openai_base_url = "https://api.x.ai/v1"
+    elif model_source == "openai":
+        openai_base_url = "https://api.openai.com/v1"
+    else:
+        openai_base_url = None  # Not an OpenAI-compatible provider
+
+    def make_image_prompt() -> str:
+        """Build prompt for image mode - direct category extraction."""
+        return (
+            f"Identify {categories_per_chunk} {specificity} categories of content found in this image. "
+            f"The image is: {image_description}. "
+            f"{'Research context: ' + research_question if research_question else ''}\n\n"
+            f"Number your categories from 1 through {categories_per_chunk} and provide concise labels only (no descriptions)."
+        )
+
+    def make_describe_prompt() -> str:
+        """Build prompt for 'both' mode - describe image content."""
+        return (
+            f"Describe the content of this image in detail. "
+            f"Include all visual elements, text, objects, people, and any other content. "
+            f"The image is: {image_description}. "
+            f"{'Research context: ' + research_question if research_question else ''}\n\n"
+            f"Provide a comprehensive text description that captures both visual and textual content."
+        )
+
+    def make_text_prompt(text_blob: str) -> str:
+        """Build prompt for extracting categories from text descriptions."""
+        return (
+            f"Identify {categories_per_chunk} {specificity} categories of content found in this description. "
+            f"The content is: {image_description}. "
+            f"{'Research context: ' + research_question + '. ' if research_question else ''}"
+            f"The description is contained within triple backticks: ```{text_blob}``` "
+            f"Number your categories from 1 through {categories_per_chunk} and provide concise labels only (no descriptions)."
+        )
+
+    def call_model_with_image(img_path, prompt_text, max_retries=6):
+        """Send an image to the model and get category extraction."""
+        encoded, ext, is_valid = _encode_image(img_path)
+        if not is_valid:
+            return None
+
+        for attempt in range(max_retries):
+            try:
+                if model_source in ["openai", "huggingface", "huggingface-together", "xai"]:
+                    endpoint = f"{openai_base_url}/chat/completions"
+                    headers = {
+                        "Content-Type": "application/json",
+                        "Authorization": f"Bearer {api_key}"
+                    }
+                    messages = [{
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": prompt_text},
+                            {"type": "image_url", "image_url": {"url": f"data:image/{ext};base64,{encoded}"}}
+                        ]
+                    }]
+                    payload = {"model": user_model, "messages": messages}
+                    if creativity is not None:
+                        payload["temperature"] = creativity
+                    response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                    response.raise_for_status()
+                    result = response.json()
+                    return result["choices"][0]["message"]["content"]
+
+                elif model_source == "anthropic":
+                    endpoint = "https://api.anthropic.com/v1/messages"
+                    headers = {
+                        "Content-Type": "application/json",
+                        "x-api-key": api_key,
+                        "anthropic-version": "2023-06-01"
+                    }
+                    media_type = f"image/{ext}" if ext else "image/jpeg"
+                    content = [
+                        {"type": "text", "text": prompt_text},
+                        {"type": "image", "source": {"type": "base64", "media_type": media_type, "data": encoded}}
+                    ]
+                    payload = {
+                        "model": user_model,
+                        "max_tokens": 2048,
+                        "messages": [{"role": "user", "content": content}],
+                    }
+                    if creativity is not None:
+                        payload["temperature"] = creativity
+                    response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                    response.raise_for_status()
+                    result = response.json()
+                    resp_content = result.get("content", [])
+                    if resp_content and resp_content[0].get("type") == "text":
+                        return resp_content[0].get("text", "")
+                    return None
+
+                elif model_source == "google":
+                    url = f"https://generativelanguage.googleapis.com/v1beta/models/{user_model}:generateContent"
+                    headers = {"x-goog-api-key": api_key, "Content-Type": "application/json"}
+                    mime_type = f"image/{ext}" if ext else "image/jpeg"
+                    parts = [
+                        {"text": prompt_text},
+                        {"inline_data": {"mime_type": mime_type, "data": encoded}}
+                    ]
+                    payload = {
+                        "contents": [{"parts": parts}],
+                        "generationConfig": {**({"temperature": creativity} if creativity is not None else {})}
+                    }
+                    response = req.post(url, headers=headers, json=payload, timeout=120)
+                    response.raise_for_status()
+                    result = response.json()
+                    if "candidates" in result and result["candidates"]:
+                        return result["candidates"][0]["content"]["parts"][0]["text"]
+                    return None
+
+                elif model_source == "mistral":
+                    endpoint = "https://api.mistral.ai/v1/chat/completions"
+                    headers = {
+                        "Content-Type": "application/json",
+                        "Authorization": f"Bearer {api_key}"
+                    }
+                    messages = [{
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": prompt_text},
+                            {"type": "image_url", "image_url": {"url": f"data:image/{ext};base64,{encoded}"}}
+                        ]
+                    }]
+                    payload = {"model": user_model, "messages": messages}
+                    if creativity is not None:
+                        payload["temperature"] = creativity
+                    response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                    response.raise_for_status()
+                    result = response.json()
+                    return result["choices"][0]["message"]["content"]
+
+            except Exception as e:
+                delay = 2 ** attempt
+                if attempt < max_retries - 1:
+                    print(f"Error processing image {img_path}: {e}. Retrying in {delay}s... (attempt {attempt + 1}/{max_retries})")
+                    import time as _time
+                    _time.sleep(delay)
+                else:
+                    print(f"Error processing image {img_path}: {e}. All {max_retries} attempts failed.")
+                    return None
+
+    def describe_image_with_vision(img_path, max_retries=6):
+        """Use vision model to describe an image's content as text."""
+        encoded, ext, is_valid = _encode_image(img_path)
+        if not is_valid:
+            return None
+        prompt_text = make_describe_prompt()
+
+        for attempt in range(max_retries):
+            try:
+                if model_source in ["openai", "huggingface", "huggingface-together", "xai"]:
+                    endpoint = f"{openai_base_url}/chat/completions"
+                    headers = {
+                        "Content-Type": "application/json",
+                        "Authorization": f"Bearer {api_key}"
+                    }
+                    messages = [{
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": prompt_text},
+                            {"type": "image_url", "image_url": {"url": f"data:image/{ext};base64,{encoded}"}}
+                        ]
+                    }]
+                    payload = {"model": user_model, "messages": messages}
+                    if creativity is not None:
+                        payload["temperature"] = creativity
+                    response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                    response.raise_for_status()
+                    result = response.json()
+                    return result["choices"][0]["message"]["content"]
+
+                elif model_source == "anthropic":
+                    endpoint = "https://api.anthropic.com/v1/messages"
+                    headers = {
+                        "Content-Type": "application/json",
+                        "x-api-key": api_key,
+                        "anthropic-version": "2023-06-01"
+                    }
+                    media_type = f"image/{ext}" if ext else "image/jpeg"
+                    content = [
+                        {"type": "text", "text": prompt_text},
+                        {"type": "image", "source": {"type": "base64", "media_type": media_type, "data": encoded}}
+                    ]
+                    payload = {
+                        "model": user_model,
+                        "max_tokens": 4096,
+                        "messages": [{"role": "user", "content": content}],
+                    }
+                    if creativity is not None:
+                        payload["temperature"] = creativity
+                    response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                    response.raise_for_status()
+                    result = response.json()
+                    resp_content = result.get("content", [])
+                    if resp_content and resp_content[0].get("type") == "text":
+                        return resp_content[0].get("text", "")
+                    return None
+
+                elif model_source == "google":
+                    url = f"https://generativelanguage.googleapis.com/v1beta/models/{user_model}:generateContent"
+                    headers = {"x-goog-api-key": api_key, "Content-Type": "application/json"}
+                    mime_type = f"image/{ext}" if ext else "image/jpeg"
+                    parts = [
+                        {"text": prompt_text},
+                        {"inline_data": {"mime_type": mime_type, "data": encoded}}
+                    ]
+                    payload = {
+                        "contents": [{"parts": parts}],
+                        "generationConfig": {**({"temperature": creativity} if creativity is not None else {})}
+                    }
+                    response = req.post(url, headers=headers, json=payload, timeout=120)
+                    response.raise_for_status()
+                    result = response.json()
+                    if "candidates" in result and result["candidates"]:
+                        return result["candidates"][0]["content"]["parts"][0]["text"]
+                    return None
+
+                elif model_source == "mistral":
+                    endpoint = "https://api.mistral.ai/v1/chat/completions"
+                    headers = {
+                        "Content-Type": "application/json",
+                        "Authorization": f"Bearer {api_key}"
+                    }
+                    messages = [{
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": prompt_text},
+                            {"type": "image_url", "image_url": {"url": f"data:image/{ext};base64,{encoded}"}}
+                        ]
+                    }]
+                    payload = {"model": user_model, "messages": messages}
+                    if creativity is not None:
+                        payload["temperature"] = creativity
+                    response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                    response.raise_for_status()
+                    result = response.json()
+                    return result["choices"][0]["message"]["content"]
+
+            except Exception as e:
+                delay = 2 ** attempt
+                if attempt < max_retries - 1:
+                    print(f"Error describing image {img_path}: {e}. Retrying in {delay}s... (attempt {attempt + 1}/{max_retries})")
+                    import time as _time
+                    _time.sleep(delay)
+                else:
+                    print(f"Error describing image {img_path}: {e}. All {max_retries} attempts failed.")
+                    return None
+
+    def call_model_with_text(prompt_text):
+        """Send text to the model for category extraction."""
+        try:
+            if model_source in ["openai", "huggingface", "huggingface-together", "xai"]:
+                endpoint = f"{openai_base_url}/chat/completions"
+                headers = {
+                    "Content-Type": "application/json",
+                    "Authorization": f"Bearer {api_key}"
+                }
+                payload = {"model": user_model, "messages": [{"role": "user", "content": prompt_text}]}
+                if creativity is not None:
+                    payload["temperature"] = creativity
+                response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                response.raise_for_status()
+                result = response.json()
+                return result["choices"][0]["message"]["content"]
+
+            elif model_source == "anthropic":
+                endpoint = "https://api.anthropic.com/v1/messages"
+                headers = {
+                    "Content-Type": "application/json",
+                    "x-api-key": api_key,
+                    "anthropic-version": "2023-06-01"
+                }
+                payload = {
+                    "model": user_model,
+                    "max_tokens": 2048,
+                    "messages": [{"role": "user", "content": prompt_text}],
+                }
+                if creativity is not None:
+                    payload["temperature"] = creativity
+                response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                response.raise_for_status()
+                result = response.json()
+                resp_content = result.get("content", [])
+                if resp_content and resp_content[0].get("type") == "text":
+                    return resp_content[0].get("text", "")
+                return None
+
+            elif model_source == "google":
+                url = f"https://generativelanguage.googleapis.com/v1beta/models/{user_model}:generateContent"
+                headers = {"x-goog-api-key": api_key, "Content-Type": "application/json"}
+                payload = {
+                    "contents": [{"parts": [{"text": prompt_text}]}],
+                    "generationConfig": {**({"temperature": creativity} if creativity is not None else {})}
+                }
+                response = req.post(url, headers=headers, json=payload, timeout=120)
+                response.raise_for_status()
+                result = response.json()
+                if "candidates" in result and result["candidates"]:
+                    return result["candidates"][0]["content"]["parts"][0]["text"]
+                return None
+
+            elif model_source == "mistral":
+                endpoint = "https://api.mistral.ai/v1/chat/completions"
+                headers = {
+                    "Content-Type": "application/json",
+                    "Authorization": f"Bearer {api_key}"
+                }
+                payload = {"model": user_model, "messages": [{"role": "user", "content": prompt_text}]}
+                if creativity is not None:
+                    payload["temperature"] = creativity
+                response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+                response.raise_for_status()
+                result = response.json()
+                return result["choices"][0]["message"]["content"]
+
+        except Exception as e:
+            print(f"Error in text mode: {e}")
+            return None
+
+    # Parse numbered list pattern
+    line_pat = re.compile(r"^\s*\d+\s*[\.\)\-]\s*(.+)$")
+
+    all_items = []
+
+    # Calculate total steps for progress tracking: (iterations * divisions) + 1 for final merge
+    total_steps = (iterations * divisions) + 1
+    current_step = 0
+
+    for pass_idx in range(iterations):
+        # Sample images for this pass
+        image_indices = list(range(n))
+        rng.shuffle(image_indices)
+
+        # Create chunks
+        chunks = [image_indices[i:i + chunk_size] for i in range(0, len(image_indices), chunk_size)][:divisions]
+
+        for chunk_idx, chunk in enumerate(tqdm(chunks, desc=f"Processing chunks (pass {pass_idx+1}/{iterations})")):
+            if not chunk:
+                continue
+
+            # Sample one random image from the full pool
+            random_idx = rng.choice(image_indices)
+            img_path = image_files[random_idx]
+
+            if mode == "image":
+                # IMAGE MODE: Send image directly for category extraction
+                prompt = make_image_prompt()
+                reply = call_model_with_image(img_path, prompt)
+
+            elif mode == "both":
+                # BOTH MODE: Describe image first, then extract categories from description
+                image_description_text = describe_image_with_vision(img_path)
+                if not image_description_text:
+                    continue
+
+                prompt = make_text_prompt(image_description_text)
+                reply = call_model_with_text(prompt)
+
+            else:
+                raise ValueError(f"Invalid mode: {mode}. Must be 'image' or 'both'.")
+
+            if reply:
+                # Extract numbered items
+                items = []
+                for raw_line in reply.splitlines():
+                    m = line_pat.match(raw_line.strip())
+                    if m:
+                        items.append(m.group(1).strip())
+                # Fallback for unnumbered lines
+                if not items:
+                    for raw_line in reply.splitlines():
+                        s = raw_line.strip()
+                        if s:
+                            items.append(s)
+                all_items.extend(items)
+
+            # Progress callback
+            current_step += 1
+            if progress_callback:
+                progress_callback(current_step, total_steps, f"Pass {pass_idx+1}/{iterations}, chunk {chunk_idx+1}/{len(chunks)}")
+
+    # Normalize and count
+    def normalize_category(cat):
+        terms = sorted([t.strip().lower() for t in str(cat).split("/")])
+        return "/".join(terms)
+
+    flat_list = [str(x).strip() for x in all_items if str(x).strip()]
+    if not flat_list:
+        raise ValueError("No categories were extracted from the images.")
+
+    df = pd.DataFrame(flat_list, columns=["Category"])
+    df["normalized"] = df["Category"].map(normalize_category)
+
+    result = (
+        df.groupby("normalized")
+          .agg(Category=("Category", lambda x: x.value_counts().index[0]),
+               counts=("Category", "size"))
+          .sort_values("counts", ascending=False)
+          .reset_index(drop=True)
+    )
+
+    # Second-pass semantic merge
+    seed_list = result["Category"].head(max_categories * 3).tolist()
+
+    second_prompt = f"""
+You are a data analyst reviewing categorized image data.
+
+Task: From the provided categories, identify and return the top {max_categories} CONCEPTUALLY UNIQUE categories.
+
+Critical Instructions:
+1) Exact duplicates are already removed.
+2) Merge SEMANTIC duplicates (same concept, different wording).
+3) When merging:
+   - Combine frequencies mentally
+   - Keep the most frequent OR clearest label
+   - Each concept appears ONLY ONCE
+4) Keep category names {specificity}.
+5) Return ONLY a numbered list of {max_categories} categories. No extra text.
+
+Pre-processed Categories (sorted by frequency, top sample):
+{seed_list}
+
+Output:
+1. category
+2. category
+...
+{max_categories}. category
+""".strip()
+
+    try:
+        if model_source in ["openai", "huggingface", "huggingface-together", "xai"]:
+            endpoint = f"{openai_base_url}/chat/completions"
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {api_key}"
+            }
+            payload = {"model": user_model, "messages": [{"role": "user", "content": second_prompt}]}
+            if creativity is not None:
+                payload["temperature"] = creativity
+            response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+            response.raise_for_status()
+            result = response.json()
+            top_categories_text = result["choices"][0]["message"]["content"]
+        elif model_source == "anthropic":
+            endpoint = "https://api.anthropic.com/v1/messages"
+            headers = {
+                "Content-Type": "application/json",
+                "x-api-key": api_key,
+                "anthropic-version": "2023-06-01"
+            }
+            payload = {
+                "model": user_model,
+                "max_tokens": 2048,
+                "messages": [{"role": "user", "content": second_prompt}],
+            }
+            if creativity is not None:
+                payload["temperature"] = creativity
+            response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+            response.raise_for_status()
+            result = response.json()
+            resp_content = result.get("content", [])
+            if resp_content and resp_content[0].get("type") == "text":
+                top_categories_text = resp_content[0].get("text", "")
+            else:
+                top_categories_text = ""
+        elif model_source == "google":
+            url = f"https://generativelanguage.googleapis.com/v1beta/models/{user_model}:generateContent"
+            headers = {"x-goog-api-key": api_key, "Content-Type": "application/json"}
+            payload = {
+                "contents": [{"parts": [{"text": second_prompt}]}],
+                "generationConfig": {**({"temperature": creativity} if creativity is not None else {})}
+            }
+            response = req.post(url, headers=headers, json=payload, timeout=120)
+            response.raise_for_status()
+            res = response.json()
+            top_categories_text = res["candidates"][0]["content"]["parts"][0]["text"]
+        elif model_source == "mistral":
+            endpoint = "https://api.mistral.ai/v1/chat/completions"
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {api_key}"
+            }
+            payload = {"model": user_model, "messages": [{"role": "user", "content": second_prompt}]}
+            if creativity is not None:
+                payload["temperature"] = creativity
+            response = req.post(endpoint, headers=headers, json=payload, timeout=120)
+            response.raise_for_status()
+            result = response.json()
+            top_categories_text = result["choices"][0]["message"]["content"]
+    except Exception as e:
+        print(f"Error in second-pass merge: {e}")
+        top_categories_text = ""
+
+    # Final progress callback for the merge step
+    if progress_callback:
+        progress_callback(total_steps, total_steps, "Merging categories")
+
+    # Parse final list
+    final = []
+    for line in top_categories_text.splitlines():
+        m = line_pat.match(line.strip())
+        if m:
+            final.append(m.group(1).strip())
+    if not final:
+        final = [l.strip("-*• ").strip() for l in top_categories_text.splitlines() if l.strip()]
+
+    print("\nTop categories:\n" + "\n".join(f"{i+1}. {c}" for i, c in enumerate(final[:max_categories])))
+
+    if filename:
+        result.to_csv(filename, index=False)
+
+    return {
+        "counts_df": result,
+        "top_categories": final[:max_categories],
+        "raw_top_text": top_categories_text
+    }