workbench 0.8.162__py3-none-any.whl → 0.8.202__py3-none-any.whl

workbench/utils/cloudwatch_handler.py

@@ -5,7 +5,7 @@ import getpass
 import time  # For managing send intervals
 
 # Workbench imports
-from workbench.utils.execution_environment import (
+from workbench_bridges.utils.execution_environment import (
     running_on_lambda,
     running_on_glue,
     running_on_ecs,

workbench/utils/cloudwatch_utils.py

@@ -0,0 +1,137 @@
+"""AWS CloudWatch utility functions for Workbench."""
+
+import time
+import logging
+from datetime import datetime, timezone
+from typing import List, Optional, Dict, Generator
+from urllib.parse import quote
+from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
+
+log = logging.getLogger("workbench")
+
+
+def get_cloudwatch_client():
+    """Get the CloudWatch Logs client using the Workbench assumed role session."""
+    session = AWSAccountClamp().boto3_session
+    return session.client("logs")
+
+
+def get_cloudwatch_logs_url(log_group: str, log_stream: str) -> Optional[str]:
+    """
+    Generate CloudWatch logs URL for the specified log group and stream.
+
+    Args:
+        log_group: Log group name (e.g., '/aws/batch/job')
+        log_stream: Log stream name
+
+    Returns:
+        CloudWatch console URL or None if unable to generate
+    """
+    try:
+        region = AWSAccountClamp().region
+
+        # URL encode the log group and stream
+        encoded_group = quote(log_group, safe="")
+        encoded_stream = quote(log_stream, safe="")
+
+        return (
+            f"https://{region}.console.aws.amazon.com/cloudwatch/home?"
+            f"region={region}#logsV2:log-groups/log-group/{encoded_group}"
+            f"/log-events/{encoded_stream}"
+        )
+    except Exception as e:  # noqa: BLE001
+        log.warning(f"Failed to generate CloudWatch logs URL: {e}")
+        return None
+
+
+def get_active_log_streams(
+    log_group_name: str, start_time_ms: int, stream_filter: Optional[str] = None, client=None
+) -> List[str]:
+    """Retrieve log streams that have events after the specified start time."""
+    if not client:
+        client = get_cloudwatch_client()
+    active_streams = []
+    stream_params = {
+        "logGroupName": log_group_name,
+        "orderBy": "LastEventTime",
+        "descending": True,
+    }
+    while True:
+        response = client.describe_log_streams(**stream_params)
+        log_streams = response.get("logStreams", [])
+        for log_stream in log_streams:
+            log_stream_name = log_stream["logStreamName"]
+            last_event_timestamp = log_stream.get("lastEventTimestamp", 0)
+            if last_event_timestamp >= start_time_ms:
+                active_streams.append(log_stream_name)
+            else:
+                break
+        if "nextToken" in response:
+            stream_params["nextToken"] = response["nextToken"]
+        else:
+            break
+    # Sort and filter streams
+    active_streams.sort()
+    if stream_filter and active_streams:
+        active_streams = [stream for stream in active_streams if stream_filter in stream]
+    return active_streams
+
+
+def stream_log_events(
+    log_group_name: str,
+    log_stream_name: str,
+    start_time: Optional[datetime] = None,
+    end_time: Optional[datetime] = None,
+    follow: bool = False,
+    client=None,
+) -> Generator[Dict, None, None]:
+    """
+    Stream log events from a specific log stream.
+    Yields:
+        Log events as dictionaries
+    """
+    if not client:
+        client = get_cloudwatch_client()
+    params = {"logGroupName": log_group_name, "logStreamName": log_stream_name, "startFromHead": True}
+    if start_time:
+        params["startTime"] = int(start_time.timestamp() * 1000)
+    if end_time:
+        params["endTime"] = int(end_time.timestamp() * 1000)
+    next_token = None
+    while True:
+        if next_token:
+            params["nextToken"] = next_token
+            params.pop("startTime", None)
+        try:
+            response = client.get_log_events(**params)
+            events = response.get("events", [])
+            for event in events:
+                event["logStreamName"] = log_stream_name
+                yield event
+            next_token = response.get("nextForwardToken")
+            # Break if no more events or same token
+            if not next_token or next_token == params.get("nextToken"):
+                if not follow:
+                    break
+                time.sleep(2)
+        except client.exceptions.ResourceNotFoundException:
+            if not follow:
+                break
+            time.sleep(2)
+
+
+def print_log_event(
+    event: dict, show_stream: bool = True, local_time: bool = True, custom_format: Optional[str] = None
+):
+    """Print a formatted log event."""
+    timestamp = datetime.fromtimestamp(event["timestamp"] / 1000, tz=timezone.utc)
+    if local_time:
+        timestamp = timestamp.astimezone()
+    message = event["message"].rstrip()
+    if custom_format:
+        # Allow custom formatting
+        print(custom_format.format(stream=event.get("logStreamName", ""), time=timestamp, message=message))
+    elif show_stream and "logStreamName" in event:
+        print(f"[{event['logStreamName']}] [{timestamp:%Y-%m-%d %I:%M%p}] {message}")
+    else:
+        print(f"[{timestamp:%H:%M:%S}] {message}")

workbench/utils/config_manager.py

@@ -4,15 +4,12 @@ import os
 import sys
 import platform
 import logging
-import importlib.resources as resources  # noqa: F401 Python 3.9 compatibility
 from typing import Any, Dict
+from importlib.resources import files, as_file
 
 # Workbench imports
 from workbench.utils.license_manager import LicenseManager
-from workbench.utils.execution_environment import running_as_service
-
-# Python 3.9 compatibility
-from workbench.utils.resource_utils import get_resource_path
+from workbench_bridges.utils.execution_environment import running_as_service
 
 
 class FatalConfigError(Exception):
@@ -172,8 +169,7 @@ class ConfigManager:
         Returns:
             str: The open source API key.
         """
-        # Python 3.9 compatibility
-        with get_resource_path("workbench.resources", "open_source_api.key") as open_source_key_path:
+        with as_file(files("workbench.resources").joinpath("open_source_api.key")) as open_source_key_path:
             with open(open_source_key_path, "r") as key_file:
                 return key_file.read().strip()
 

workbench/utils/endpoint_utils.py

@@ -7,9 +7,7 @@ from typing import Union, Optional
 import pandas as pd
 
 # Workbench Imports
-from workbench.api.feature_set import FeatureSet
-from workbench.api.model import Model
-from workbench.api.endpoint import Endpoint
+from workbench.api import FeatureSet, Model, Endpoint
 
 # Set up the log
 log = logging.getLogger("workbench")
@@ -77,7 +75,7 @@ def internal_model_data_url(endpoint_config_name: str, session: boto3.Session) -
         return None
 
 
-def fs_training_data(end: Endpoint) -> pd.DataFrame:
+def get_training_data(end: Endpoint) -> pd.DataFrame:
     """Code to get the training data from the FeatureSet used to train the Model
 
     Args:
@@ -100,7 +98,7 @@ def fs_training_data(end: Endpoint) -> pd.DataFrame:
     return train_df
 
 
-def fs_evaluation_data(end: Endpoint) -> pd.DataFrame:
+def get_evaluation_data(end: Endpoint) -> pd.DataFrame:
     """Code to get the evaluation data from the FeatureSet NOT used for training
 
     Args:
@@ -178,11 +176,11 @@ if __name__ == "__main__":
     print(model_data_url)
 
     # Get the training data
-    my_train_df = fs_training_data(my_endpoint)
+    my_train_df = get_training_data(my_endpoint)
     print(my_train_df)
 
     # Get the evaluation data
-    my_eval_df = fs_evaluation_data(my_endpoint)
+    my_eval_df = get_evaluation_data(my_endpoint)
     print(my_eval_df)
 
     # Backtrack to the FeatureSet

workbench/utils/license_manager.py

@@ -6,15 +6,12 @@ import json
 import logging
 import requests
 from typing import Union
-import importlib.resources as resources  # noqa: F401 Python 3.9 compatibility
 from datetime import datetime
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.asymmetric import padding
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.backends import default_backend
-
-# Python 3.9 compatibility
-from workbench.utils.resource_utils import get_resource_path
+from importlib.resources import files, as_file
 
 
 class FatalLicenseError(Exception):
@@ -140,8 +137,7 @@ class LicenseManager:
         Returns:
             The public key as an object.
         """
-        # Python 3.9 compatibility
-        with get_resource_path("workbench.resources", "signature_verify_pub.pem") as public_key_path:
+        with as_file(files("workbench.resources").joinpath("signature_verify_pub.pem")) as public_key_path:
             with open(public_key_path, "rb") as key_file:
                 public_key_data = key_file.read()
 

workbench/utils/model_utils.py

@@ -3,6 +3,7 @@
 import logging
 import pandas as pd
 import numpy as np
+from scipy.stats import spearmanr
 import importlib.resources
 from pathlib import Path
 import os
@@ -92,6 +93,38 @@ def get_custom_script_path(package: str, script_name: str) -> Path:
     return script_path
 
 
+def proximity_model_local(model: "Model"):
+    """Create a Proximity Model for this Model
+
+    Args:
+        model (Model): The Model/FeatureSet used to create the proximity model
+
+    Returns:
+        Proximity: The proximity model
+    """
+    from workbench.algorithms.dataframe.proximity import Proximity  # noqa: F401 (avoid circular import)
+    from workbench.api import Model, FeatureSet  # noqa: F401 (avoid circular import)
+
+    # Get Feature and Target Columns from the existing given Model
+    features = model.features()
+    target = model.target()
+
+    # Backtrack our FeatureSet to get the ID column
+    fs = FeatureSet(model.get_input())
+    id_column = fs.id_column
+
+    # Create the Proximity Model from both the full FeatureSet and the Model training data
+    full_df = fs.pull_dataframe()
+    model_df = model.training_view().pull_dataframe()
+
+    # Mark rows that are in the model
+    model_ids = set(model_df[id_column])
+    full_df["in_model"] = full_df[id_column].isin(model_ids)
+
+    # Create and return the Proximity Model
+    return Proximity(full_df, id_column, features, target, track_columns=features)
+
+
 def proximity_model(model: "Model", prox_model_name: str, track_columns: list = None) -> "Model":
     """Create a proximity model based on the given model
 
@@ -139,9 +172,6 @@ def uq_model(model: "Model", uq_model_name: str, train_all_data: bool = False) -
     """
     from workbench.api import Model, ModelType, FeatureSet  # noqa: F401 (avoid circular import)
 
-    # Get the custom script path for the UQ model
-    script_path = get_custom_script_path("uq_models", "meta_uq.template")
-
     # Get Feature and Target Columns from the existing given Model
     features = model.features()
     target = model.target()
@@ -156,12 +186,25 @@ def uq_model(model: "Model", uq_model_name: str, train_all_data: bool = False) -
         description=f"UQ Model for {model.name}",
         tags=["uq", model.name],
         train_all_data=train_all_data,
-        custom_script=script_path,
         custom_args={"id_column": fs.id_column, "track_columns": [target]},
     )
     return uq_model
 
 
+def safe_extract_tarfile(tar_path: str, extract_path: str) -> None:
+    """
+    Extract a tarball safely, using data filter if available.
+
+    The filter parameter was backported to Python 3.8+, 3.9+, 3.10.13+, 3.11+
+    as a security patch, but may not be present in older patch versions.
+    """
+    with tarfile.open(tar_path, "r:gz") as tar:
+        if hasattr(tarfile, "data_filter"):
+            tar.extractall(path=extract_path, filter="data")
+        else:
+            tar.extractall(path=extract_path)
+
+
 def load_category_mappings_from_s3(model_artifact_uri: str) -> Optional[dict]:
     """
     Download and extract category mappings from a model artifact in S3.
@@ -180,8 +223,7 @@ def load_category_mappings_from_s3(model_artifact_uri: str) -> Optional[dict]:
         wr.s3.download(path=model_artifact_uri, local_file=local_tar_path)
 
         # Extract tarball
-        with tarfile.open(local_tar_path, "r:gz") as tar:
-            tar.extractall(path=tmpdir, filter="data")
+        safe_extract_tarfile(local_tar_path, tmpdir)
 
         # Look for category mappings in base directory only
         mappings_path = os.path.join(tmpdir, "category_mappings.json")
@@ -220,28 +262,41 @@ def uq_metrics(df: pd.DataFrame, target_col: str) -> Dict[str, Any]:
     # --- Coverage and Interval Width ---
     if "q_025" in df.columns and "q_975" in df.columns:
         lower_95, upper_95 = df["q_025"], df["q_975"]
+        lower_90, upper_90 = df["q_05"], df["q_95"]
+        lower_80, upper_80 = df["q_10"], df["q_90"]
+        lower_68 = df.get("q_16", df["q_10"])  # fallback to 80% interval
+        upper_68 = df.get("q_84", df["q_90"])  # fallback to 80% interval
         lower_50, upper_50 = df["q_25"], df["q_75"]
     elif "prediction_std" in df.columns:
         lower_95 = df["prediction"] - 1.96 * df["prediction_std"]
         upper_95 = df["prediction"] + 1.96 * df["prediction_std"]
+        lower_90 = df["prediction"] - 1.645 * df["prediction_std"]
+        upper_90 = df["prediction"] + 1.645 * df["prediction_std"]
+        lower_80 = df["prediction"] - 1.282 * df["prediction_std"]
+        upper_80 = df["prediction"] + 1.282 * df["prediction_std"]
+        lower_68 = df["prediction"] - 1.0 * df["prediction_std"]
+        upper_68 = df["prediction"] + 1.0 * df["prediction_std"]
         lower_50 = df["prediction"] - 0.674 * df["prediction_std"]
         upper_50 = df["prediction"] + 0.674 * df["prediction_std"]
     else:
         raise ValueError(
             "Either quantile columns (q_025, q_975, q_25, q_75) or 'prediction_std' column must be present."
         )
+    median_std = df["prediction_std"].median()
     coverage_95 = np.mean((df[target_col] >= lower_95) & (df[target_col] <= upper_95))
-    coverage_50 = np.mean((df[target_col] >= lower_50) & (df[target_col] <= upper_50))
-    avg_width_95 = np.mean(upper_95 - lower_95)
-    avg_width_50 = np.mean(upper_50 - lower_50)
+    coverage_90 = np.mean((df[target_col] >= lower_90) & (df[target_col] <= upper_90))
+    coverage_80 = np.mean((df[target_col] >= lower_80) & (df[target_col] <= upper_80))
+    coverage_68 = np.mean((df[target_col] >= lower_68) & (df[target_col] <= upper_68))
+    median_width_95 = np.median(upper_95 - lower_95)
+    median_width_90 = np.median(upper_90 - lower_90)
+    median_width_80 = np.median(upper_80 - lower_80)
+    median_width_50 = np.median(upper_50 - lower_50)
+    median_width_68 = np.median(upper_68 - lower_68)
 
     # --- CRPS (measures calibration + sharpness) ---
-    if "prediction_std" in df.columns:
-        z = (df[target_col] - df["prediction"]) / df["prediction_std"]
-        crps = df["prediction_std"] * (z * (2 * norm.cdf(z) - 1) + 2 * norm.pdf(z) - 1 / np.sqrt(np.pi))
-        mean_crps = np.mean(crps)
-    else:
-        mean_crps = np.nan
+    z = (df[target_col] - df["prediction"]) / df["prediction_std"]
+    crps = df["prediction_std"] * (z * (2 * norm.cdf(z) - 1) + 2 * norm.pdf(z) - 1 / np.sqrt(np.pi))
+    mean_crps = np.mean(crps)
 
     # --- Interval Score @ 95% (penalizes miscoverage) ---
     alpha_95 = 0.05
@@ -252,31 +307,43 @@ def uq_metrics(df: pd.DataFrame, target_col: str) -> Dict[str, Any]:
     )
     mean_is_95 = np.mean(is_95)
 
-    # --- Adaptive Calibration (correlation between errors and uncertainty) ---
+    # --- Interval to Error Correlation ---
     abs_residuals = np.abs(df[target_col] - df["prediction"])
-    width_95 = upper_95 - lower_95
-    adaptive_calibration = np.corrcoef(abs_residuals, width_95)[0, 1]
+    width_68 = upper_68 - lower_68
+
+    # Spearman correlation for robustness
+    interval_to_error_corr = spearmanr(width_68, abs_residuals)[0]
 
     # Collect results
     results = {
+        "coverage_68": coverage_68,
+        "coverage_80": coverage_80,
+        "coverage_90": coverage_90,
         "coverage_95": coverage_95,
-        "coverage_50": coverage_50,
-        "avg_width_95": avg_width_95,
-        "avg_width_50": avg_width_50,
-        "crps": mean_crps,
-        "interval_score_95": mean_is_95,
-        "adaptive_calibration": adaptive_calibration,
+        "median_std": median_std,
+        "median_width_50": median_width_50,
+        "median_width_68": median_width_68,
+        "median_width_80": median_width_80,
+        "median_width_90": median_width_90,
+        "median_width_95": median_width_95,
+        "interval_to_error_corr": interval_to_error_corr,
         "n_samples": len(df),
     }
 
     print("\n=== UQ Metrics ===")
+    print(f"Coverage @ 68%: {coverage_68:.3f} (target: 0.68)")
+    print(f"Coverage @ 80%: {coverage_80:.3f} (target: 0.80)")
+    print(f"Coverage @ 90%: {coverage_90:.3f} (target: 0.90)")
     print(f"Coverage @ 95%: {coverage_95:.3f} (target: 0.95)")
-    print(f"Coverage @ 50%: {coverage_50:.3f} (target: 0.50)")
-    print(f"Average 95% Width: {avg_width_95:.3f}")
-    print(f"Average 50% Width: {avg_width_50:.3f}")
+    print(f"Median Prediction StdDev: {median_std:.3f}")
+    print(f"Median 50% Width: {median_width_50:.3f}")
+    print(f"Median 68% Width: {median_width_68:.3f}")
+    print(f"Median 80% Width: {median_width_80:.3f}")
+    print(f"Median 90% Width: {median_width_90:.3f}")
+    print(f"Median 95% Width: {median_width_95:.3f}")
     print(f"CRPS: {mean_crps:.3f} (lower is better)")
     print(f"Interval Score 95%: {mean_is_95:.3f} (lower is better)")
-    print(f"Adaptive Calibration: {adaptive_calibration:.3f} (higher is better, target: >0.5)")
+    print(f"Interval/Error Corr: {interval_to_error_corr:.3f} (higher is better, target: >0.5)")
     print(f"Samples: {len(df)}")
     return results
 
@@ -313,9 +380,3 @@ if __name__ == "__main__":
     df = end.auto_inference(capture=True)
     results = uq_metrics(df, target_col="solubility")
     print(results)
-
-    # Test the uq_metrics function
-    end = Endpoint("aqsol-uq-100")
-    df = end.auto_inference(capture=True)
-    results = uq_metrics(df, target_col="solubility")
-    print(results)

workbench/utils/monitor_utils.py

@@ -14,7 +14,7 @@ from workbench.utils.s3_utils import read_content_from_s3
 log = logging.getLogger("workbench")
 
 
-def pull_data_capture(data_capture_path, max_files=1) -> Union[pd.DataFrame, None]:
+def pull_data_capture_for_testing(data_capture_path, max_files=1) -> Union[pd.DataFrame, None]:
     """
     Read and process captured data from S3.
 
@@ -26,7 +26,12 @@ def pull_data_capture(data_capture_path, max_files=1) -> Union[pd.DataFrame, Non
 
     Returns:
         Union[pd.DataFrame, None]: A dataframe of the captured data (or None if no data is found).
+
+    Notes:
+        This method is really only for testing and debugging.
     """
+    log.important("This method is for testing and debugging only.")
+
     # List files in the specified S3 path
     files = wr.s3.list_objects(data_capture_path)
     if not files:
@@ -64,59 +69,53 @@ def pull_data_capture(data_capture_path, max_files=1) -> Union[pd.DataFrame, Non
 def process_data_capture(df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
     """
     Process the captured data DataFrame to extract input and output data.
-    Continues processing even if individual files are malformed.
+    Handles cases where input or output might not be captured.
+
     Args:
         df (DataFrame): DataFrame with captured data.
     Returns:
         tuple[DataFrame, DataFrame]: Input and output DataFrames.
     """
+
+    def parse_endpoint_data(data: dict) -> pd.DataFrame:
+        """Parse endpoint data based on encoding type."""
+        encoding = data["encoding"].upper()
+
+        if encoding == "CSV":
+            return pd.read_csv(StringIO(data["data"]))
+        elif encoding == "JSON":
+            json_data = json.loads(data["data"])
+            if isinstance(json_data, dict):
+                return pd.DataFrame({k: [v] if not isinstance(v, list) else v for k, v in json_data.items()})
+            else:
+                return pd.DataFrame(json_data)
+        else:
+            return None  # Unknown encoding
+
     input_dfs = []
     output_dfs = []
 
-    for idx, row in df.iterrows():
+    # Use itertuples() instead of iterrows() for better performance
+    for row in df.itertuples(index=True):
         try:
-            capture_data = row["captureData"]
-
-            # Check if this capture has the required fields (all or nothing)
-            if "endpointInput" not in capture_data:
-                log.warning(f"Row {idx}: No endpointInput found in capture data.")
-                continue
-
-            if "endpointOutput" not in capture_data:
-                log.critical(
-                    f"Row {idx}: No endpointOutput found in capture data. DataCapture needs to include Output capture!"
-                )
-                continue
-
-            # Process input data
-            input_data = capture_data["endpointInput"]
-            if input_data["encoding"].upper() == "CSV":
-                input_df = pd.read_csv(StringIO(input_data["data"]))
-            elif input_data["encoding"].upper() == "JSON":
-                json_data = json.loads(input_data["data"])
-                if isinstance(json_data, dict):
-                    input_df = pd.DataFrame({k: [v] if not isinstance(v, list) else v for k, v in json_data.items()})
-                else:
-                    input_df = pd.DataFrame(json_data)
-
-            # Process output data
-            output_data = capture_data["endpointOutput"]
-            if output_data["encoding"].upper() == "CSV":
-                output_df = pd.read_csv(StringIO(output_data["data"]))
-            elif output_data["encoding"].upper() == "JSON":
-                json_data = json.loads(output_data["data"])
-                if isinstance(json_data, dict):
-                    output_df = pd.DataFrame({k: [v] if not isinstance(v, list) else v for k, v in json_data.items()})
-                else:
-                    output_df = pd.DataFrame(json_data)
-
-            # If we get here, both processed successfully
-            input_dfs.append(input_df)
-            output_dfs.append(output_df)
+            capture_data = row.captureData
+
+            # Process input data if present
+            if "endpointInput" in capture_data:
+                input_df = parse_endpoint_data(capture_data["endpointInput"])
+                if input_df is not None:
+                    input_dfs.append(input_df)
+
+            # Process output data if present
+            if "endpointOutput" in capture_data:
+                output_df = parse_endpoint_data(capture_data["endpointOutput"])
+                if output_df is not None:
+                    output_dfs.append(output_df)
 
         except Exception as e:
-            log.error(f"Row {idx}: Failed to process row: {e}")
+            log.debug(f"Row {row.Index}: Failed to process row: {e}")
             continue
+
     # Combine and return results
     return (
         pd.concat(input_dfs, ignore_index=True) if input_dfs else pd.DataFrame(),
@@ -178,23 +177,6 @@ def parse_monitoring_results(results_json: str) -> Dict[str, Any]:
         return {"error": str(e)}
 
 
-"""TEMP
-                # If the status is "CompletedWithViolations", we grab the lastest
-                # violation file and add it to the result
-                if status == "CompletedWithViolations":
-                    violation_file = f"{self.monitoring_path}/
-                    {last_run['CreationTime'].strftime('%Y/%m/%d')}/constraint_violations.json"
-                    if wr.s3.does_object_exist(violation_file):
-                        violations_json = read_content_from_s3(violation_file)
-                        violations = parse_monitoring_results(violations_json)
-                        result["violations"] = violations.get("constraint_violations", [])
-                        result["violation_count"] = len(result["violations"])
-                    else:
-                        result["violations"] = []
-                        result["violation_count"] = 0
-"""
-
-
 def preprocessing_script(feature_list: list[str]) -> str:
     """
     A preprocessing script for monitoring jobs.
@@ -245,8 +227,8 @@ if __name__ == "__main__":
     from workbench.api.monitor import Monitor
 
     # Test pulling data capture
-    mon = Monitor("caco2-pappab-class-0")
-    df = pull_data_capture(mon.data_capture_path)
+    mon = Monitor("abalone-regression-rt")
+    df = pull_data_capture_for_testing(mon.data_capture_path)
     print("Data Capture:")
     print(df.head())
 
@@ -262,4 +244,4 @@ if __name__ == "__main__":
     # Test preprocessing script
     script = preprocessing_script(["feature1", "feature2", "feature3"])
     print("\nPreprocessing Script:")
-    print(script)
+    # print(script)

workbench/utils/pandas_utils.py

@@ -152,7 +152,7 @@ def compare_dataframes(df1: pd.DataFrame, df2: pd.DataFrame, display_columns: li
 
     # Check for differences in common columns
     for column in common_columns:
-        if pd.api.types.is_string_dtype(df1[column]) or pd.api.types.is_string_dtype(df2[column]):
+        if pd.api.types.is_string_dtype(df1[column]) and pd.api.types.is_string_dtype(df2[column]):
             # String comparison with NaNs treated as equal
             differences = ~(df1[column].fillna("") == df2[column].fillna(""))
         elif pd.api.types.is_float_dtype(df1[column]) or pd.api.types.is_float_dtype(df2[column]):
@@ -161,8 +161,8 @@ def compare_dataframes(df1: pd.DataFrame, df2: pd.DataFrame, display_columns: li
                 pd.isna(df1[column]) & pd.isna(df2[column])
             )
         else:
-            # Other types (e.g., int) with NaNs treated as equal
-            differences = ~(df1[column].fillna(0) == df2[column].fillna(0))
+            # Other types (int, Int64, etc.) - compare with NaNs treated as equal
+            differences = (df1[column] != df2[column]) & ~(pd.isna(df1[column]) & pd.isna(df2[column]))
 
         # If differences exist, display them
         if differences.any():