linkml-store 0.1.13__py3-none-any.whl → 0.2.0__py3-none-any.whl

linkml_store/utils/format_utils.py

@@ -27,6 +27,7 @@ class Format(Enum):
     JSON = "json"
     JSONL = "jsonl"
     YAML = "yaml"
+    YAMLL = "yamll"
     TSV = "tsv"
     CSV = "csv"
     PYTHON = "python"
@@ -46,6 +47,7 @@ class Format(Enum):
             ".jsonl": cls.JSONL,
             ".yaml": cls.YAML,
             ".yml": cls.YAML,
+            ".yamll": cls.YAMLL,
             ".tsv": cls.TSV,
             ".csv": cls.CSV,
             ".py": cls.PYTHON,
@@ -63,6 +65,9 @@ class Format(Enum):
     def is_dump_format(self):
         return self in [Format.SQLDUMP_DUCKDB, Format.SQLDUMP_POSTGRES, Format.DUMP_MONGODB]
 
+    def is_xsv(self):
+        return self in [Format.TSV, Format.CSV]
+
 
 def load_objects_from_url(
     url: str,
@@ -94,6 +99,9 @@ def process_file(
     """
     Process a single file and return a list of objects.
     """
+    if format == Format.YAMLL:
+        format = Format.YAML
+        expected_type = list
     if format == Format.JSON:
         objs = json.load(f)
     elif format == Format.JSONL:
@@ -101,6 +109,8 @@ def process_file(
     elif format == Format.YAML:
         if expected_type and expected_type == list:  # noqa E721
             objs = list(yaml.safe_load_all(f))
+            # allow YAML with a `---` with no object before it
+            objs = [obj for obj in objs if obj is not None]
         else:
             objs = yaml.safe_load(f)
     elif format in [Format.TSV, Format.CSV]:
@@ -135,11 +145,14 @@ def load_objects(
     compression: Optional[str] = None,
     expected_type: Optional[Type] = None,
     header_comment_token: Optional[str] = None,
+    select_query: Optional[str] = None,
 ) -> List[Dict[str, Any]]:
     """
     Load objects from a file or archive in supported formats.
     For tgz archives, it processes all files and concatenates the results.
 
+    TODO: Add schema hints for CSV/TSV parsing.
+
     :param file_path: The path to the file or archive.
     :param format: The format of the file. Can be a Format enum or a string value.
     :param compression: The compression type. Supports 'gz' for gzip and 'tgz' for tar.gz.
@@ -177,98 +190,22 @@ def load_objects(
             all_objects = process_file(f, format, expected_type, header_comment_token)
 
     logger.debug(f"Loaded {len(all_objects)} objects from {file_path}")
+    if select_query:
+        import jsonpath_ng as jp
+
+        path_expr = jp.parse(select_query)
+        new_objs = []
+        for obj in all_objects:
+            for match in path_expr.find(obj):
+                logging.debug(f"Match: {match.value}")
+                if isinstance(match.value, list):
+                    new_objs.extend(match.value)
+                else:
+                    new_objs.append(match.value)
+        all_objects = new_objs
     return all_objects
 
 
-def xxxload_objects(
-    file_path: Union[str, Path],
-    format: Union[Format, str] = None,
-    compression: Optional[str] = None,
-    expected_type: Type = None,
-    header_comment_token: Optional[str] = None,
-) -> List[Dict[str, Any]]:
-    """
-    Load objects from a file in JSON, JSONLines, YAML, CSV, or TSV format.
-
-    >>> load_objects("tests/input/test_data/data.csv")
-    [{'id': '1', 'name': 'John', 'age': '30'},
-     {'id': '2', 'name': 'Alice', 'age': '25'}, {'id': '3', 'name': 'Bob', 'age': '35'}]
-
-    :param file_path: The path to the file.
-    :param format: The format of the file. Can be a Format enum or a string value.
-    :param expected_type: The target type to load the objects into, e.g. list
-    :return: A list of dictionaries representing the loaded objects.
-    """
-    if isinstance(format, str):
-        format = Format(format)
-
-    if isinstance(file_path, Path):
-        file_path = str(file_path)
-
-    if not format and (file_path.endswith(".parquet") or file_path.endswith(".pq")):
-        format = Format.PARQUET
-    if not format and file_path.endswith(".tsv"):
-        format = Format.TSV
-    if not format and file_path.endswith(".csv"):
-        format = Format.CSV
-    if not format and file_path.endswith(".py"):
-        format = Format.PYTHON
-
-    mode = "r"
-    if format == Format.PARQUET:
-        mode = "rb"
-
-    if file_path == "-":
-        # set file_path to be a stream from stdin
-        f = sys.stdin
-    else:
-        f = open(file_path, mode)
-
-    if format == Format.JSON or (not format and file_path.endswith(".json")):
-        objs = json.load(f)
-    elif format == Format.JSONL or (not format and file_path.endswith(".jsonl")):
-        objs = [json.loads(line) for line in f]
-    elif format == Format.YAML or (not format and (file_path.endswith(".yaml") or file_path.endswith(".yml"))):
-        if expected_type and expected_type == list:  # noqa E721
-            objs = list(yaml.safe_load_all(f))
-        else:
-            objs = yaml.safe_load(f)
-    elif format == Format.TSV or format == Format.CSV:
-        # Skip initial comment lines if comment_char is set
-        if header_comment_token:
-            # Store the original position
-            original_pos = f.tell()
-
-            # Read and store lines until we find a non-comment line
-            lines = []
-            for line in f:
-                if not line.startswith(header_comment_token):
-                    break
-                lines.append(line)
-
-            # Go back to the original position
-            f.seek(original_pos)
-
-            # Skip the comment lines we found
-            for _ in lines:
-                f.readline()
-        if format == Format.TSV:
-            reader = csv.DictReader(f, delimiter="\t")
-        else:
-            reader = csv.DictReader(f)
-        objs = list(reader)
-    elif format == Format.PARQUET:
-        import pyarrow.parquet as pq
-
-        table = pq.read_table(f)
-        objs = table.to_pandas().to_dict(orient="records")
-    else:
-        raise ValueError(f"Unsupported file format: {file_path}")
-    if not isinstance(objs, list):
-        objs = [objs]
-    return objs
-
-
 def write_output(
     data: Union[List[Dict[str, Any]], Dict[str, Any], pd.DataFrame],
     format: Union[Format, str] = Format.YAML,
@@ -329,7 +266,7 @@ def render_output(
     if format == Format.FORMATTED:
         if not isinstance(data, pd.DataFrame):
             data = pd.DataFrame(data)
-        return str(data)
+        return data.to_string(max_rows=None)
 
     if isinstance(data, pd.DataFrame):
         data = data.to_dict(orient="records")

linkml_store/utils/llm_utils.py

@@ -0,0 +1,96 @@
+from typing import Callable, List, Optional
+
+from tiktoken import Encoding
+
+MODEL_TOKEN_MAPPING = {
+    "gpt-4o-mini": 128_000,
+    "gpt-4o": 128_000,
+    "gpt-4o-2024-05-13": 128_000,
+    "gpt-4": 8192,
+    "gpt-4-0314": 8192,
+    "gpt-4-0613": 8192,
+    "gpt-4-32k": 32768,
+    "gpt-4-32k-0314": 32768,
+    "gpt-4-32k-0613": 32768,
+    "gpt-3.5-turbo": 4096,
+    "gpt-3.5-turbo-0301": 4096,
+    "gpt-3.5-turbo-0613": 4096,
+    "gpt-3.5-turbo-16k": 16385,
+    "gpt-3.5-turbo-16k-0613": 16385,
+    "gpt-3.5-turbo-instruct": 4096,
+    "text-ada-001": 2049,
+    "ada": 2049,
+    "text-babbage-001": 2040,
+    "babbage": 2049,
+    "text-curie-001": 2049,
+    "curie": 2049,
+    "davinci": 2049,
+    "text-davinci-003": 4097,
+    "text-davinci-002": 4097,
+    "code-davinci-002": 8001,
+    "code-davinci-001": 8001,
+    "code-cushman-002": 2048,
+    "code-cushman-001": 2048,
+    "claude": 200_000,
+    "llama-3": 200_000,
+}
+
+
+def render_formatted_text(
+    render_func: Callable,
+    values: List[str],
+    encoding: Encoding,
+    token_limit: int,
+    additional_text: Optional[str] = None,
+) -> str:
+    """
+    Render a formatted text string with a given object, encoding, and token limit.
+
+    >>> from tiktoken import encoding_for_model
+    >>> encoding = encoding_for_model("gpt-4o-mini")
+    >>> names = ["Alice", "Bob", "DoctorHippopotamusMcHippopotamusFace"]
+    >>> f = lambda x: f"Hello, {' '.join(x)}!"
+    >>> render_formatted_text(f, names, encoding, 4096)
+    'Hello, Alice Bob DoctorHippopotamusMcHippopotamusFace!'
+    >>> render_formatted_text(f, names, encoding, 5)
+    'Hello, Alice Bob!'
+
+    :param render_func: Rendering function
+    :param values: Values to render
+    :param encoding: Encoding
+    :param token_limit: Token limit
+    :param additional_text: Additional text to consider
+    :return:
+    """
+    text = render_func(values)
+    if additional_text:
+        token_limit -= len(encoding.encode(additional_text))
+    text_length = len(encoding.encode(text))
+    if text_length <= token_limit:
+        return text
+    if not values:
+        raise ValueError(f"Cannot fit text into token limit: {text_length} > {token_limit}")
+    return render_formatted_text(render_func, values[0:-1], encoding=encoding, token_limit=token_limit)
+
+
+def get_token_limit(model_name: str) -> int:
+    """
+    Estimate the token limit for a model.
+
+    >>> get_token_limit("gpt-4o-mini")
+    128000
+
+    also works with nested names:
+
+    >>> get_token_limit("my/claude-opus")
+    200000
+
+
+    :param model_name: Model name
+    :return: Estimated token limit
+    """
+    # sort MODEL_TOKEN_MAPPING by key length to ensure that the longest model names are checked first
+    for model, token_limit in sorted(MODEL_TOKEN_MAPPING.items(), key=lambda x: len(x[0]), reverse=True):
+        if model in model_name:
+            return token_limit
+    return 4096

linkml_store/utils/object_utils.py

@@ -1,6 +1,6 @@
 import json
 from copy import deepcopy
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel
 
@@ -29,7 +29,7 @@ def object_path_update(
     """
     if isinstance(obj, BaseModel):
         typ = type(obj)
-        obj = obj.dict()
+        obj = obj.model_dump(exclude_none=True)
         obj = object_path_update(obj, path, value)
         return typ(**obj)
     obj = deepcopy(obj)
@@ -45,6 +45,8 @@ def object_path_update(
                 obj.append({})
             obj = obj[index]
         else:
+            if part in obj and obj[part] is None:
+                del obj[part]
             obj = obj.setdefault(part, {})
     last_part = parts[-1]
     if "[" in last_part:
@@ -58,6 +60,41 @@ def object_path_update(
     return ret_obj
 
 
+def object_path_get(obj: Union[BaseModel, Dict[str, Any]], path: str, default_value=None) -> Any:
+    """
+    Retrieves a value from a nested object based on a path description. The path to the
+    desired field is given in dot and bracket notation (e.g., 'a[0].b.c[1]').
+
+    :param obj: The dictionary object to be updated.
+    :type obj: Dict[str, Any]
+    :param path: The path string indicating where to place the value within the object.
+    :type path: str
+    :return: The value at the specified path.
+    :rtype: Any
+
+    **Example**::
+
+    >>> data = {'persons': [{'foo': {'bar': 1}}]}
+    >>> object_path_get(data, 'persons[0].foo.bar')
+    1
+    >>> object_path_get(data, 'persons[0].foo')
+    {'bar': 1}
+    >>> object_path_get({}, 'not there', "NA")
+    'NA'
+    """
+    if isinstance(obj, BaseModel):
+        obj = obj.dict()
+    parts = path.split(".")
+    for part in parts:
+        if "[" in part:
+            key, index = part[:-1].split("[")
+            index = int(index)
+            obj = obj[key][index]
+        else:
+            obj = obj.get(part, default_value)
+    return obj
+
+
 def parse_update_expression(expr: str) -> Union[tuple[str, Any], None]:
     """
     Parse a string expression of the form 'path.to.field=value' into a path and a value.
@@ -79,3 +116,67 @@ def clean_empties(value: Union[Dict, List]) -> Any:
     elif isinstance(value, list):
         value = [v for v in (clean_empties(v) for v in value) if v is not None]
     return value
+
+
+def select_nested(data: dict, paths: List[Union[str, List[str]]], current_path=None) -> Optional[dict]:
+    """
+    Select nested attributes from a complex dictionary based on selector strings.
+
+    Args:
+    data (dict): The input nested dictionary.
+    selectors (list): A list of selector strings.
+
+    Returns:
+    dict: A new dictionary with the same structure, but only the selected attributes.
+
+    Example:
+    >>> data = {
+    ...     "person": {
+    ...         "name": "John Doe",
+    ...         "age": 30,
+    ...         "address": {
+    ...             "street": "123 Main St",
+    ...             "city": "Anytown",
+    ...             "country": "USA"
+    ...         },
+    ...         "phones": [
+    ...             {"type": "home", "number": "555-1234"},
+    ...             {"type": "work", "number": "555-5678"}
+    ...         ]
+    ...     },
+    ...     "company": {
+    ...         "name": "Acme Inc",
+    ...         "location": "New York"
+    ...     }
+    ... }
+    >>> select_nested(data, ["person.address.street", "person.address.city"])
+    {'person': {'address': {'street': '123 Main St', 'city': 'Anytown'}}}
+    >>> select_nested(data, ["person.phones.number", "person.phones.type"])
+    {'person': {'phones': [{'type': 'home', 'number': '555-1234'}, {'type': 'work', 'number': '555-5678'}]}}
+    >>> select_nested(data, ["person"])
+    {'person': {'name': 'John Doe', 'age': 30, 'address': {'street': '123 Main St', 'city': 'Anytown',
+     'country': 'USA'}, 'phones': [{'type': 'home', 'number': '555-1234'}, {'type': 'work', 'number': '555-5678'}]}}
+    >>> select_nested(data, ["person.phones.type"])
+    {'person': {'phones': [{'type': 'home'}, {'type': 'work'}]}}
+    """
+    if current_path is None:
+        current_path = []
+    matching_paths = []
+    for path in paths:
+        if isinstance(path, str):
+            path = path.split(".")
+        if path == current_path:
+            return data
+        if path[: len(current_path)] == current_path:
+            matching_paths.append(path)
+    if not matching_paths:
+        return None
+    if isinstance(data, dict):
+        new_obj = {k: select_nested(v, matching_paths, current_path + [k]) for k, v in data.items()}
+        new_obj = {k: v for k, v in new_obj.items() if v is not None}
+        return new_obj
+    if isinstance(data, list):
+        new_obj = [select_nested(v, matching_paths, current_path + []) for i, v in enumerate(data)]
+        new_obj = [v for v in new_obj if v is not None]
+        return new_obj
+    return data

linkml_store/utils/pandas_utils.py

@@ -1,7 +1,59 @@
-from typing import Dict, List, Tuple, Union
+import logging
+from typing import Any, Dict, List, Tuple, Union
 
 import pandas as pd
 
+logger = logging.getLogger(__name__)
+
+
+def flatten_dict(d: Dict[str, Any], parent_key: str = "", sep: str = ".") -> Dict[str, Any]:
+    """
+    Recursively flatten a nested dictionary.
+
+    Args:
+        d (Dict[str, Any]): The dictionary to flatten.
+        parent_key (str): The parent key for nested dictionaries.
+        sep (str): The separator to use between keys.
+
+    Returns:
+        Dict[str, Any]: A flattened dictionary.
+
+    >>> flatten_dict({'a': 1, 'b': {'c': 2, 'd': {'e': 3}}})
+    {'a': 1, 'b.c': 2, 'b.d.e': 3}
+    """
+    items = []
+    for k, v in d.items():
+        new_key = f"{parent_key}{sep}{k}" if parent_key else k
+        if isinstance(v, dict):
+            items.extend(flatten_dict(v, new_key, sep=sep).items())
+        else:
+            items.append((new_key, v))
+    return dict(items)
+
+
+def nested_objects_to_dataframe(data: List[Dict[str, Any]]) -> pd.DataFrame:
+    """
+    Convert a list of nested objects to a flattened pandas DataFrame.
+
+    Args:
+        data (List[Dict[str, Any]]): A list of nested dictionaries.
+
+    Returns:
+        pd.DataFrame: A flattened DataFrame.
+
+    >>> data = [
+    ...     {"person": {"name": "Alice", "age": 30}, "job": {"title": "Engineer", "salary": 75000}},
+    ...     {"person": {"name": "Bob", "age": 35}, "job": {"title": "Manager", "salary": 85000}}
+    ... ]
+    >>> df = nested_objects_to_dataframe(data)
+    >>> df.columns.tolist()
+    ['person.name', 'person.age', 'job.title', 'job.salary']
+    >>> df['person.name'].tolist()
+    ['Alice', 'Bob']
+    """
+    flattened_data = [flatten_dict(item) for item in data]
+    return pd.DataFrame(flattened_data)
+
 
 def facet_summary_to_dataframe_unmelted(
     facet_summary: Dict[Union[str, Tuple[str, ...]], List[Tuple[Union[str, Tuple[str, ...]], int]]]
@@ -22,7 +74,8 @@ def facet_summary_to_dataframe_unmelted(
                     categories, value = cat_val_tuple[:-1], cat_val_tuple[-1]
                 row = {"Value": value}
                 for i, facet in enumerate(facet_type):
-                    row[facet] = categories[i]
+                    logger.debug(f"FT={facet_type} i={i} Facet: {facet}, categories: {categories}")
+                    row[facet] = categories[i] if len(categories) > i else None
                 rows.append(row)
 
     df = pd.DataFrame(rows)

linkml_store/utils/sklearn_utils.py

@@ -0,0 +1,193 @@
+import logging
+import os
+import re
+import shutil
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+from linkml_runtime.utils.formatutils import underscore
+from sklearn.preprocessing import LabelEncoder, OneHotEncoder
+from sklearn.tree import DecisionTreeClassifier, _tree, export_graphviz
+
+logger = logging.getLogger(__name__)
+
+
+def tree_to_nested_expression(
+    tree: DecisionTreeClassifier,
+    feature_names: List[str],
+    categorical_features: Optional[List[str]] = None,
+    feature_encoders: Optional[Dict[str, Union[OneHotEncoder, LabelEncoder]]] = None,
+    target_encoder: Optional[LabelEncoder] = None,
+) -> str:
+    """
+    Convert a trained scikit-learn DecisionTreeClassifier to a nested Python conditional expression.
+
+    Args:
+    tree (DecisionTreeClassifier): A trained decision tree classifier.
+    feature_names (list): List of feature names (including one-hot encoded feature names).
+    categorical_features (list): List of original categorical feature names.
+    feature_encoders (dict): Dictionary mapping feature names to their respective OneHotEncoders or LabelEncoders.
+    target_encoder (LabelEncoder, optional): LabelEncoder for the target variable if it's categorical.
+
+    Returns:
+    str: A string representing the nested Python conditional expression.
+
+    Example:
+    >>> import numpy as np
+    >>> from sklearn.tree import DecisionTreeClassifier
+    >>> from sklearn.preprocessing import OneHotEncoder, LabelEncoder
+    >>>
+    >>> # Prepare sample data
+    >>> X = np.array([[0, 'A'], [0, 'B'], [1, 'A'], [1, 'B']])
+    >>> y = np.array(['No', 'Yes', 'Yes', 'No'])
+    >>>
+    >>> # Prepare the encoders
+    >>> feature_encoders = {'feature2': OneHotEncoder(sparse_output=False, handle_unknown='ignore')}
+    >>> target_encoder = LabelEncoder()
+    >>>
+    >>> # Encode the categorical feature and target
+    >>> X_encoded = np.column_stack([
+    ...     X[:, 0],
+    ...     feature_encoders['feature2'].fit_transform(X[:, 1].reshape(-1, 1))
+    ... ])
+    >>> y_encoded = target_encoder.fit_transform(y)
+    >>>
+    >>> # Train the decision tree
+    >>> clf = DecisionTreeClassifier(random_state=42)
+    >>> clf.fit(X_encoded, y_encoded)
+    DecisionTreeClassifier(random_state=42)
+    >>>
+    >>> # Convert to nested expression
+    >>> feature_names = ['feature1', 'feature2_A', 'feature2_B']
+    >>> categorical_features = ['feature2']
+    >>> expression = tree_to_nested_expression(clf, feature_names,
+    ...                                        categorical_features, feature_encoders, target_encoder)
+    >>> print(expression)
+    (("Yes" if ({feature1} <= 0.5000) else "No") if ({feature2} == "A")
+       else ("No" if ({feature1} <= 0.5000) else "Yes"))
+    """
+    tree_ = tree.tree_
+    feature_name = [feature_names[i] if i != _tree.TREE_UNDEFINED else "undefined!" for i in tree_.feature]
+
+    categorical_features = set(categorical_features or [])
+
+    def get_original_feature_name(name):
+        return name.split("_")[0] if "_" in name else name
+
+    def recurse(node):
+        if tree_.feature[node] != _tree.TREE_UNDEFINED:
+            name = feature_name[node]
+            threshold = tree_.threshold[node]
+            original_name = get_original_feature_name(name)
+            original_name_safe = underscore(original_name)
+            name_safe = underscore(name)
+
+            original_name_safe = "{" + original_name_safe + "}"
+            name_safe = "{" + name_safe + "}"
+
+            if original_name in categorical_features:
+                if feature_encoders is None or original_name not in feature_encoders:
+                    raise ValueError(f"Encoder is required for categorical feature {original_name}")
+
+                encoder = feature_encoders[original_name]
+                if isinstance(encoder, OneHotEncoder):
+                    # For one-hot encoded features, we check if the specific category is present
+                    category = name.split("_", 1)[1]  # Get everything after the first underscore
+                    condition = f'{original_name_safe} == "{category}"'
+                elif isinstance(encoder, LabelEncoder):
+                    category = encoder.inverse_transform([int(threshold)])[0]
+                    condition = f'{original_name_safe} == "{category}"'
+                else:
+                    raise ValueError(f"Unsupported encoder type for feature {original_name}")
+            else:
+                if np.isinf(threshold):
+                    condition = "True"
+                else:
+                    condition = f"{name_safe} <= {threshold:.4f}"
+
+            left_expr = recurse(tree_.children_left[node])
+            right_expr = recurse(tree_.children_right[node])
+
+            return f"({left_expr} if ({condition}) else {right_expr})"
+        else:
+            class_index = np.argmax(tree_.value[node])
+            if target_encoder:
+                class_label = target_encoder.inverse_transform([class_index])[0]
+                return f'"{class_label}"'
+            else:
+                return str(class_index)
+
+    return recurse(0)
+
+
+def escape_label(s: str) -> str:
+    """Escape special characters in label strings."""
+    s = str(s)
+    return re.sub(r"([<>])", r"\\\1", s)
+
+
+def visualize_decision_tree(
+    clf: DecisionTreeClassifier,
+    feature_names: List[str],
+    class_names: List[str] = None,
+    output_file: Union[Path, str] = "decision_tree.png",
+) -> None:
+    """
+    Generate a visualization of the decision tree and save it as a PNG file.
+
+    :param clf: Trained DecisionTreeClassifier
+    :param feature_names: List of feature names
+    :param class_names: List of class names (optional)
+    :param output_file: The name of the file to save the visualization (default: "decision_tree.png")
+
+    >>> # Create a sample dataset
+    >>> import pandas as pd
+    >>> data = pd.DataFrame({
+    ...     'age': [25, 30, 35, 40, 45],
+    ...     'income': [50000, 60000, 70000, 80000, 90000],
+    ...     'credit_score': [600, 650, 700, 750, 800],
+    ...     'approved': ['No', 'No', 'Yes', 'Yes', 'Yes']
+    ... })
+    >>>
+    >>> # Prepare features and target
+    >>> X = data[['age', 'income', 'credit_score']]
+    >>> y = data['approved']
+    >>>
+    >>> # Encode target variable
+    >>> le = LabelEncoder()
+    >>> y_encoded = le.fit_transform(y)
+    >>>
+    >>> # Train a decision tree
+    >>> clf = DecisionTreeClassifier(random_state=42)
+    >>> _ = clf.fit(X, y_encoded)
+    >>> # Visualize the tree
+    >>> visualize_decision_tree(clf, X.columns.tolist(), le.classes_, "tests/output/test_tree.png")
+    """
+    # Escape special characters in feature names and class names
+    escaped_feature_names = [escape_label(name) for name in feature_names]
+    escaped_class_names = [escape_label(name) for name in (class_names if class_names is not None else [])]
+
+    import graphviz
+
+    dot_data = export_graphviz(
+        clf,
+        out_file=None,
+        feature_names=escaped_feature_names,
+        class_names=escaped_class_names,
+        filled=True,
+        rounded=True,
+        special_characters=True,
+    )
+    # dot_data = escape_label(dot_data)
+    logger.info(f"Dot: {dot_data}")
+    dot_path = shutil.which("dot")
+    if not dot_path:
+        logger.warning("Graphviz 'dot' executable not found in PATH. Skipping visualization.")
+        return
+    os.environ["GRAPHVIZ_DOT"] = dot_path
+
+    graph = graphviz.Source(dot_data)
+    if isinstance(output_file, Path):
+        output_file = str(output_file)
+    graph.render(output_file.rsplit(".", 1)[0], format="png", cleanup=True)