classifier-toolkit 0.3.6__tar.gz → 0.5.0__tar.gz

{classifier_toolkit-0.3.6 → classifier_toolkit-0.5.0}/.gitignore

@@ -178,4 +178,8 @@ cython_debug/
 catboost_info/
 
 # Gitkeep
-!/**/.gitkeep
+!/**/.gitkeep
+
+# Claude Code local state (worktrees, personal settings)
+.claude/worktrees/
+.claude/settings.local.json

{classifier_toolkit-0.3.6 → classifier_toolkit-0.5.0}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.5
 Name: classifier-toolkit
-Version: 0.3.6
+Version: 0.5.0
 Project-URL: Documentation, https://supreme-adventure-jg5qkyr.pages.github.io/
 Author-email: "senih.yilmaz" <senih.yilmaz@qonto.com>, "jeremy.fraoua" <jeremy.fraoua@qonto.com>, "gauthier.marquand" <gauthier.marquand@qonto.com>, "arnaud.alepee" <arnaud.alepee@qonto.com>
 License-File: LICENSE
-Requires-Python: <3.14,>=3.9
+Requires-Python: <3.14,>=3.10
 Requires-Dist: catboost<2.0.0,>=1.2.2
 Requires-Dist: category-encoders<3.0.0,>=2.6.3
 Requires-Dist: colorama<0.5.0,>=0.4.6
@@ -18,6 +18,7 @@ Requires-Dist: polars<2.0.0,>=1.2.1
 Requires-Dist: pyarrow>=18.0.0
 Requires-Dist: scikit-learn<2.0.0,>=1.4.0
 Requires-Dist: scipy>=1.11.0
+Requires-Dist: seaborn<0.14.0,>=0.13.0
 Requires-Dist: shap>=0.46.0
 Requires-Dist: statsmodels<0.15.0,>=0.14.2
 Requires-Dist: tabulate<0.10.0,>=0.9.0
@@ -60,13 +61,17 @@ This library is published in the PyPI directory. To install, users can run pip i
 
 ### Usage
 
-This library automates binary classification tasks in the finance domain, specifically for default and fraud labeling. It includes several packages designed to address the main steps in any machine learning/data science task:
+This library automates binary and multiclass classification workflows. It is independent of the modelled problem: the class of interest is configured through `pos_label` (the positive class for binary targets, the class of interest for multiclass ones). It includes several packages designed to address the main steps in any machine learning/data science task:
 
-1. **EDA**: accessible via `EDA_Toolkit`. Provides EDA and feature engineering functionality with all necessary visualizations.
+1. **EDA**: accessible via `EDAToolkit`. Provides EDA and feature engineering functionality with all necessary visualizations.
 2. **Feature Reduction**: filter-style pre-selection pipeline (expert rules, low variance, drift, predictive power, counter-intuitive direction, high correlation).
 3. **Feature Selection**: wrapper and embedded methods (RFE, Boruta, Sequential, Bayesian, ElasticNet, MetaSelector).
-4. **Model Training**: accessible via `Tuner`. Hyperparameter optimization (grid search, Bayesian via Optuna) for LightGBM and CatBoost with train/val/test evaluation.
-5. Evaluation and reporting: To be implemented.
+4. **Model Training**: accessible via `Tuner`. Hyperparameter optimization (grid search, Bayesian via Optuna) for LightGBM and CatBoost with train/val/test evaluation, including multiclass objectives and class weights.
+5. **Risk Class**: `RiskClassBuilderDP` / `RiskClassBuilder` turn model scores into risk classes with statistically validated, ordered event rates.
+6. **Model Validation**: `Evaluator` computes metrics, calibration and stability checks for any model exposing `predict_proba`.
+7. **Calibration**: `OVRHistogramCalibrator` and `plot_reliability_curves` for One-vs-Rest probability calibration.
+8. **Explainability**: SHAP-based explanations, interactions and misclassification diagnostics.
+9. **Data Partition**: temporal-aware train/test splitting, preprocessing and dtype optimization.
 
 For detailed usage, refer to the documentation.
 
@@ -139,16 +144,46 @@ For detailed usage, refer to the documentation.
 - **Model Training**: Hyperparameter optimization for LightGBM and CatBoost, with support for grid search and Bayesian optimization (via Optuna).
 
   ```python
-  from classifier_toolkit.model_training.hyper_parameter_tuning import Tuner
+  from classifier_toolkit.model_training.hyper_parameter_tuning.tuner import Tuner
+
+  tuner = Tuner(
+      X=X_train, y=y_train,
+      model_name="lightgbm",
+      X_val=X_val, y_val=y_val,
+      X_test=X_test, y_test=y_test,
+      search_method="bayesian",
+      n_trials=50,
+      optimization_metric="prauc",
+  )
+  result = tuner.tune()
+
+  best_model = result["best_model"]
+  result["trials_results"]  # full trial results (DataFrame)
+  ```
+
+  Reported metrics are `auc`, `prauc`, `ks`, `log_loss` and `brier` (`ks`/`brier` are binary-only). Custom parameter search spaces can be defined via `ModelParams` and `ParamRange`.
+
+- **Risk Class**: Builds risk classes from model scores. `RiskClassBuilderDP` searches bin edges with dynamic programming so that each class' observed event rate falls in a target band (`target_ranges`, required), then checks that adjacent classes are statistically distinguishable; `RiskClassBuilder` discovers classes with KMeans.
 
-  tuner = Tuner(model_name="lightgbm", metric="auc", search_type="bayesian", n_trials=50)
-  tuner.fit(X_train, y_train, X_val, y_val, X_test=X_test, y_test=y_test)
+  ```python
+  from classifier_toolkit.risk_class import RiskClassBuilderDP
 
-  best_model = tuner.best_model
-  tuner.results_df  # full trial results
+  builder = RiskClassBuilderDP(
+      target_col="target",
+      target_ranges=[(0.00, 0.02), (0.02, 0.05), (0.05, 0.10)],
+      min_obs_per_bin=200,
+  )
+  result = builder.build(train_proba, y_train)
+  print(result.bins, result.n_classes)
   ```
 
-  Custom parameter search spaces can be defined via `ModelParams` and `ParamRange`.
+- **Model Validation**: `Evaluator` evaluates any model with `predict_proba` (metrics, ROC/PR/calibration/threshold plots), `score_distribution_psi` and `compare_score_distributions` compare a reference score distribution with a current one, and `evaluate_risk_classes_dp` checks calibration within each risk class.
+
+- **Calibration**: `OVRHistogramCalibrator` calibrates multiclass (or binary) probabilities One-vs-Rest with histogram binning; `plot_reliability_curves` shows raw vs calibrated reliability per class.
+
+- **Explainability**: `TreeSHAPExplainer`, `ExplainabilityToolkit`, SHAP plots, pairwise interaction analysis and a `MisclassificationAnalyzer` that explains confusion-matrix quadrants (per-class SHAP for multiclass models).
+
+- **Data Partition**: Temporal-aware train/test splitting, preprocessing helpers and dtype optimization.
 
 ### Development & CI/CD
 
@@ -162,8 +197,8 @@ This project uses modern tooling for fast and efficient development workflows:
 #### CI/CD Pipeline
 Our CI/CD pipeline is optimized for speed and efficiency:
 
-- **Parallel Test Execution**: Tests are split into two groups (`eda` and `feature_selection`) that run simultaneously, reducing test time by ~50%
-- **Shared Caching**: Both parallel jobs share the same dependency cache (~1.7GB), avoiding duplicate downloads
+- **Parallel Test Execution**: One test job per directory under `tests/` (discovered automatically, so new test groups are picked up without editing the workflow), all running simultaneously
+- **Shared Caching**: The parallel jobs share the same dependency cache, avoiding duplicate downloads
 - **Smart Test Reruns**: Failed tests run first (`pytest --lf --ff`) for faster feedback on fixes
 - **Master Protection**: Build tests only run on `master` branch and PRs targeting `master`, saving CI resources on feature branches
 - **Automatic Linting**: Code quality checks (Ruff, SQLFluff) run on every push
@@ -188,5 +223,5 @@ uv build
 
 ### Future Work
 The next planned improvements and additions to the library include:
-* Developing comprehensive evaluation and reporting tools to assist with model assessment.
+* Extending the evaluation and reporting tools in `model_validation` (more checks and report formats).
 * Expanding documentation to include architecture diagrams and detailed usage examples.

{classifier_toolkit-0.3.6 → classifier_toolkit-0.5.0}/README.md

@@ -29,13 +29,17 @@ This library is published in the PyPI directory. To install, users can run pip i
 
 ### Usage
 
-This library automates binary classification tasks in the finance domain, specifically for default and fraud labeling. It includes several packages designed to address the main steps in any machine learning/data science task:
+This library automates binary and multiclass classification workflows. It is independent of the modelled problem: the class of interest is configured through `pos_label` (the positive class for binary targets, the class of interest for multiclass ones). It includes several packages designed to address the main steps in any machine learning/data science task:
 
-1. **EDA**: accessible via `EDA_Toolkit`. Provides EDA and feature engineering functionality with all necessary visualizations.
+1. **EDA**: accessible via `EDAToolkit`. Provides EDA and feature engineering functionality with all necessary visualizations.
 2. **Feature Reduction**: filter-style pre-selection pipeline (expert rules, low variance, drift, predictive power, counter-intuitive direction, high correlation).
 3. **Feature Selection**: wrapper and embedded methods (RFE, Boruta, Sequential, Bayesian, ElasticNet, MetaSelector).
-4. **Model Training**: accessible via `Tuner`. Hyperparameter optimization (grid search, Bayesian via Optuna) for LightGBM and CatBoost with train/val/test evaluation.
-5. Evaluation and reporting: To be implemented.
+4. **Model Training**: accessible via `Tuner`. Hyperparameter optimization (grid search, Bayesian via Optuna) for LightGBM and CatBoost with train/val/test evaluation, including multiclass objectives and class weights.
+5. **Risk Class**: `RiskClassBuilderDP` / `RiskClassBuilder` turn model scores into risk classes with statistically validated, ordered event rates.
+6. **Model Validation**: `Evaluator` computes metrics, calibration and stability checks for any model exposing `predict_proba`.
+7. **Calibration**: `OVRHistogramCalibrator` and `plot_reliability_curves` for One-vs-Rest probability calibration.
+8. **Explainability**: SHAP-based explanations, interactions and misclassification diagnostics.
+9. **Data Partition**: temporal-aware train/test splitting, preprocessing and dtype optimization.
 
 For detailed usage, refer to the documentation.
 
@@ -108,16 +112,46 @@ For detailed usage, refer to the documentation.
 - **Model Training**: Hyperparameter optimization for LightGBM and CatBoost, with support for grid search and Bayesian optimization (via Optuna).
 
   ```python
-  from classifier_toolkit.model_training.hyper_parameter_tuning import Tuner
+  from classifier_toolkit.model_training.hyper_parameter_tuning.tuner import Tuner
+
+  tuner = Tuner(
+      X=X_train, y=y_train,
+      model_name="lightgbm",
+      X_val=X_val, y_val=y_val,
+      X_test=X_test, y_test=y_test,
+      search_method="bayesian",
+      n_trials=50,
+      optimization_metric="prauc",
+  )
+  result = tuner.tune()
+
+  best_model = result["best_model"]
+  result["trials_results"]  # full trial results (DataFrame)
+  ```
+
+  Reported metrics are `auc`, `prauc`, `ks`, `log_loss` and `brier` (`ks`/`brier` are binary-only). Custom parameter search spaces can be defined via `ModelParams` and `ParamRange`.
+
+- **Risk Class**: Builds risk classes from model scores. `RiskClassBuilderDP` searches bin edges with dynamic programming so that each class' observed event rate falls in a target band (`target_ranges`, required), then checks that adjacent classes are statistically distinguishable; `RiskClassBuilder` discovers classes with KMeans.
 
-  tuner = Tuner(model_name="lightgbm", metric="auc", search_type="bayesian", n_trials=50)
-  tuner.fit(X_train, y_train, X_val, y_val, X_test=X_test, y_test=y_test)
+  ```python
+  from classifier_toolkit.risk_class import RiskClassBuilderDP
 
-  best_model = tuner.best_model
-  tuner.results_df  # full trial results
+  builder = RiskClassBuilderDP(
+      target_col="target",
+      target_ranges=[(0.00, 0.02), (0.02, 0.05), (0.05, 0.10)],
+      min_obs_per_bin=200,
+  )
+  result = builder.build(train_proba, y_train)
+  print(result.bins, result.n_classes)
   ```
 
-  Custom parameter search spaces can be defined via `ModelParams` and `ParamRange`.
+- **Model Validation**: `Evaluator` evaluates any model with `predict_proba` (metrics, ROC/PR/calibration/threshold plots), `score_distribution_psi` and `compare_score_distributions` compare a reference score distribution with a current one, and `evaluate_risk_classes_dp` checks calibration within each risk class.
+
+- **Calibration**: `OVRHistogramCalibrator` calibrates multiclass (or binary) probabilities One-vs-Rest with histogram binning; `plot_reliability_curves` shows raw vs calibrated reliability per class.
+
+- **Explainability**: `TreeSHAPExplainer`, `ExplainabilityToolkit`, SHAP plots, pairwise interaction analysis and a `MisclassificationAnalyzer` that explains confusion-matrix quadrants (per-class SHAP for multiclass models).
+
+- **Data Partition**: Temporal-aware train/test splitting, preprocessing helpers and dtype optimization.
 
 ### Development & CI/CD
 
@@ -131,8 +165,8 @@ This project uses modern tooling for fast and efficient development workflows:
 #### CI/CD Pipeline
 Our CI/CD pipeline is optimized for speed and efficiency:
 
-- **Parallel Test Execution**: Tests are split into two groups (`eda` and `feature_selection`) that run simultaneously, reducing test time by ~50%
-- **Shared Caching**: Both parallel jobs share the same dependency cache (~1.7GB), avoiding duplicate downloads
+- **Parallel Test Execution**: One test job per directory under `tests/` (discovered automatically, so new test groups are picked up without editing the workflow), all running simultaneously
+- **Shared Caching**: The parallel jobs share the same dependency cache, avoiding duplicate downloads
 - **Smart Test Reruns**: Failed tests run first (`pytest --lf --ff`) for faster feedback on fixes
 - **Master Protection**: Build tests only run on `master` branch and PRs targeting `master`, saving CI resources on feature branches
 - **Automatic Linting**: Code quality checks (Ruff, SQLFluff) run on every push
@@ -157,5 +191,5 @@ uv build
 
 ### Future Work
 The next planned improvements and additions to the library include:
-* Developing comprehensive evaluation and reporting tools to assist with model assessment.
+* Extending the evaluation and reporting tools in `model_validation` (more checks and report formats).
 * Expanding documentation to include architecture diagrams and detailed usage examples.

{classifier_toolkit-0.3.6 → classifier_toolkit-0.5.0}/classifier_toolkit/data_partition/split_train_test.py

@@ -11,200 +11,121 @@ def split_data_recent_orgs(
     date_col="observation_date",
 ):
     """
-    Split data into train and test sets based on the most recent organizations.
-    Takes the most recent organizations (by first observation date) for the test set
-    based on the specified test_size proportion, and the remaining organizations for training.
-    Iteratively adjusts the split to maintain default rate balance within tolerance.
+    Split data into train and test sets on the most recent entities.
+
+    Entities (``org_id_col``) are ordered by their first observation date, and
+    the most recent ones go to the test set, so the test set mimics entities
+    the model will meet after deployment. The number of test entities is
+    ``test_size`` of all entities, adjusted as little as possible so that the
+    train and test event rates differ by at most ``rate_tolerance``: sizes are
+    tried from the target outwards (n, n-1, n+1, n-2, ...) within +/-30% of
+    it (at least +/-5), and the first one within tolerance is kept. If none
+    is, the size with the smallest event-rate difference is used. The target
+    is ``round(test_size * n_entities)``, at least 1.
 
     Parameters:
     -----------
     df : pandas.DataFrame
         The input dataframe
     test_size : float, default=0.2
-        Proportion of organizations to include in test set (most recent ones)
+        Share of **entities** (not rows) put in the test set. Recent entities
+        usually have fewer observations, so the test share of rows is lower.
     rate_tolerance : float, default=0.05
-        Maximum acceptable difference in default rates between train and test (5%)
-    max_iterations : int, default=50
-        Maximum number of iterations to find a balanced split
+        Maximum absolute difference between the train and test event rates
+        (``0.05`` = 5 percentage points; use a value in proportion to the
+        event rate).
     label_col : str, default='default_label'
-        Column containing the binary labels
+        Column containing the binary (0/1) target
     org_id_col : str, default='org_id'
-        Column containing organization IDs
+        Column containing the entity IDs. Entities that share a first
+        observation date are ordered by ID, so the IDs must be mutually
+        comparable (not a mix of ``int`` and ``str``).
     date_col : str, default='observation_date'
         Column containing the observation dates
 
     Returns:
     --------
     tuple
-        (train_df, test_df) - Training and testing dataframes
+        ``(train_df, test_df, stats)``, where ``stats`` is the same dict as
+        the other splitters return (``train_orgs``, ``test_orgs``,
+        ``train_records``, ``test_records``, ``train_target_rate``,
+        ``test_target_rate``, ``target_rate_diff``, ...).
     """
-    # Ensure date column is datetime
     df = df.copy()
     if not pd.api.types.is_datetime64_dtype(df[date_col]):
         df[date_col] = pd.to_datetime(df[date_col])
 
-    # Get the earliest observation date for each organization
+    # First observation date of each entity, oldest first. The entity ID
+    # breaks ties, so which entities of a same-date group end up in test is
+    # deterministic.
     org_first_dates = (
         df.groupby(org_id_col)[date_col]
         .min()
         .reset_index()
-        .sort_values(date_col, ascending=True)
+        .sort_values([date_col, org_id_col], kind="stable")
     )
-
-    # Calculate number of organizations for test set
     total_orgs = len(org_first_dates)
-    n_test_orgs = int(total_orgs * test_size)
-
-    print("=== RECENT ORGANIZATIONS SPLIT WITH RATE TOLERANCE ===")
-    print(f"Test size target: {test_size:.1%}")
-    print(f"Rate tolerance: {rate_tolerance:.1%}")
-    print(f"Total organizations: {total_orgs}")
+    if total_orgs < 2:
+        raise ValueError(
+            f"split_data_recent_orgs needs at least 2 entities in {org_id_col!r}, "
+            f"got {total_orgs}."
+        )
+    n_test_orgs = max(1, round(total_orgs * test_size))
 
-    # Try different splits around the target to find one within tolerance
-    best_split = None
-    best_rate_diff = float("inf")
+    print("=== RECENT ENTITIES SPLIT WITH RATE TOLERANCE ===")
+    print(f"Test size target: {test_size:.1%} of {total_orgs} entities")
+    print(f"Rate tolerance: {rate_tolerance:.2%}")
 
-    # Search range: ±30% of target test size
+    # Sizes from the target outwards, so the first one within tolerance is
+    # the closest to the requested test_size.
     search_range = max(5, int(n_test_orgs * 0.3))
-    min_test_orgs = max(1, n_test_orgs - search_range)
-    max_test_orgs = min(total_orgs - 1, n_test_orgs + search_range)
-
-    print(
-        f"Searching for balanced split (testing {min_test_orgs} to {max_test_orgs} orgs)..."
-    )
-
-    for test_org_count in range(min_test_orgs, max_test_orgs + 1):
-        # Take organizations for this test size
-        test_orgs_data = org_first_dates.tail(test_org_count)
-
-        # Create temporary splits
-        temp_train = df[~df[org_id_col].isin(test_orgs_data[org_id_col].values)]
-        temp_test = df[df[org_id_col].isin(test_orgs_data[org_id_col].values)]
-
-        if len(temp_train) == 0 or len(temp_test) == 0:
+    low = max(1, n_test_orgs - search_range)
+    high = min(total_orgs - 1, n_test_orgs + search_range)
+    candidates = sorted(range(low, high + 1), key=lambda n: (abs(n - n_test_orgs), n))
+
+    best_count, best_rate_diff = None, float("inf")
+    for test_org_count in candidates:
+        test_ids = org_first_dates[org_id_col].tail(test_org_count).to_numpy()
+        is_test = df[org_id_col].isin(test_ids)
+        if is_test.all() or not is_test.any():
             continue
-
-        # Calculate default rates
-        train_rate = temp_train[label_col].mean()
-        test_rate = temp_test[label_col].mean()
-        rate_diff = abs(train_rate - test_rate)
-
-        # Check if this split is better
+        rate_diff = abs(
+            df.loc[~is_test, label_col].mean() - df.loc[is_test, label_col].mean()
+        )
         if rate_diff < best_rate_diff:
-            best_rate_diff = rate_diff
-            best_split = {
-                "test_orgs_data": test_orgs_data,
-                "train_rate": train_rate,
-                "test_rate": test_rate,
-                "rate_diff": rate_diff,
-                "test_org_count": test_org_count,
-                "actual_test_size": len(temp_test) / len(df),
-            }
-
-        # Stop if we found an acceptable split
+            best_count, best_rate_diff = test_org_count, rate_diff
         if rate_diff <= rate_tolerance:
             print(
-                f"✓ Found balanced split at {test_org_count} test orgs (rate diff: {rate_diff:.4f})"
+                f"Found a balanced split at {test_org_count} test entities "
+                f"(event-rate difference: {rate_diff:.4f})"
             )
             break
+    else:
+        if best_count is None:
+            best_count = n_test_orgs
+            print("Could not find any valid split; using the target size.")
+        else:
+            print(
+                f"Best achievable event-rate difference: {best_rate_diff:.4f} "
+                f"(exceeds the tolerance {rate_tolerance:.4f}); using "
+                f"{best_count} test entities."
+            )
 
-    if best_split is None:
-        # Fallback to original target
-        best_split = {
-            "test_orgs_data": org_first_dates.tail(n_test_orgs),
-            "test_org_count": n_test_orgs,
-        }
-        print("⚠️  Could not find any valid split. Using target split.")
-    elif best_rate_diff > rate_tolerance:
-        print(
-            f"⚠️  Best achievable rate difference: {best_rate_diff:.4f} (exceeds tolerance: {rate_tolerance:.4f})"
-        )
-
-    # Use the best split found
-    test_orgs_data = best_split["test_orgs_data"]
-
-    # Split data into train and test
-    train_df = df[~df[org_id_col].isin(test_orgs_data[org_id_col].values)].copy()
-    test_df = df[df[org_id_col].isin(test_orgs_data[org_id_col].values)].copy()
-
-    # Calculate statistics
-    total_orgs = df[org_id_col].nunique()
-    total_records = len(df)
-    total_pos = df[df[label_col] == 1].shape[0]
-    total_neg = df[df[label_col] == 0].shape[0]
-
-    train_orgs = train_df[org_id_col].nunique()
-    train_records = len(train_df)
-    train_pos = train_df[train_df[label_col] == 1].shape[0]
-    train_neg = train_df[train_df[label_col] == 0].shape[0]
-
-    test_orgs = test_df[org_id_col].nunique()
-    test_records = len(test_df)
-    test_pos = test_df[test_df[label_col] == 1].shape[0]
-    test_neg = test_df[test_df[label_col] == 0].shape[0]
-
-    # Verify no overlap
-    train_org_set = set(train_df[org_id_col].unique())
-    test_org_set = set(test_df[org_id_col].unique())
-    overlap = train_org_set.intersection(test_org_set)
-
-    # Print statistics
-    print("\n=== FINAL SPLIT STATISTICS ===")
-    print(f"Test organizations: {test_orgs} ({test_orgs / total_orgs:.1%})")
-    print(f"Train organizations: {train_orgs} ({train_orgs / total_orgs:.1%})")
-
-    print(f"\nTotal records: {total_records}")
-    print(f"Train records: {train_records} ({train_records / total_records:.2%})")
-    print(f"Test records: {test_records} ({test_records / total_records:.2%})")
-
-    print("\nPositive samples:")
-    print(f"Total: {total_pos}")
-    print(f"Train: {train_pos} ({train_pos / total_pos:.2%})")
-    print(f"Test: {test_pos} ({test_pos / total_pos:.2%})")
-
-    print("\nNegative samples:")
-    print(f"Total: {total_neg}")
-    print(f"Train: {train_neg} ({train_neg / total_neg:.2%})")
-    print(f"Test: {test_neg} ({test_neg / total_neg:.2%})")
-
-    print("\nDefault rates:")
-    print(f"Overall: {df[label_col].mean():.4f}")
-    print(f"Train: {train_df[label_col].mean():.4f}")
-    print(f"Test: {test_df[label_col].mean():.4f}")
-    print(
-        f"Difference: {abs(train_df[label_col].mean() - test_df[label_col].mean()):.4f}"
-    )
-
-    print(f"\nOrganization overlap: {len(overlap)} organizations (should be 0)")
+    test_orgs_data = org_first_dates.tail(best_count)
+    is_test = df[org_id_col].isin(test_orgs_data[org_id_col].to_numpy())
+    train_df = df[~is_test].copy()
+    test_df = df[is_test].copy()
 
-    print("\nDate ranges:")
-    print(
-        "Cut-off date for splitting (earliest test org first loan):",
-        test_orgs_data[date_col].min().strftime("%Y-%m-%d"),
-    )
-    print(
-        f"Training orgs (first loan): {train_df.groupby(org_id_col)[date_col].min().min().strftime('%Y-%m-%d')} to {train_df.groupby(org_id_col)[date_col].min().max().strftime('%Y-%m-%d')}"
-    )
-    print(
-        f"Testing orgs (first loan): {test_df.groupby(org_id_col)[date_col].min().min().strftime('%Y-%m-%d')} to {test_df.groupby(org_id_col)[date_col].min().max().strftime('%Y-%m-%d')}"
-    )
-    print(
-        f"Training all observations: {train_df[date_col].min().strftime('%Y-%m-%d')} to {train_df[date_col].max().strftime('%Y-%m-%d')}"
+    stats = _calculate_split_statistics(
+        df, train_df, test_df, label_col, org_id_col, date_col
     )
+    overlap = set(train_df[org_id_col]) & set(test_df[org_id_col])
+    print(f"\nEntity overlap: {len(overlap)} (should be 0)")
     print(
-        f"Testing all observations: {test_df[date_col].min().strftime('%Y-%m-%d')} to {test_df[date_col].max().strftime('%Y-%m-%d')}"
-    )
-
-    return (
-        train_df,
-        test_df,
-        (
-            train_df.shape,
-            test_df.shape,
-            train_df[train_df[label_col] == 1].shape,
-            test_df[test_df[label_col] == 1].shape,
-        ),
+        "Cut-off (first observation of the earliest test entity): "
+        f"{test_orgs_data[date_col].min():%Y-%m-%d}"
     )
+    return train_df, test_df, stats
 
 
 def split_data_stratified_by_org(

classifier_toolkit-0.5.0/classifier_toolkit/datasets/__init__.py

@@ -0,0 +1,3 @@
+from ._demo import ID_COLS, TARGET_COL, make_demo_data
+
+__all__ = ["ID_COLS", "TARGET_COL", "make_demo_data"]