splitpilot 0.1.0__py3-none-any.whl

splitpilot/__init__.py

@@ -0,0 +1,7 @@
+from .core.pilot import Pilot
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Pilot",
+]

splitpilot/core/__init__.py

@@ -0,0 +1 @@
+"""Core SplitPilot functionality."""

splitpilot/core/pilot.py

@@ -0,0 +1,67 @@
+from .profiler import DatasetProfiler
+from .recommender import SplitRecommender
+from .splitter import DatasetSplitter
+
+
+class Pilot:
+    """
+    High-level interface for analyzing a dataset
+    and creating a leakage-aware train/test split.
+    """
+
+    def __init__(self, df, target):
+        self.df = df
+        self.target = target
+
+        self._profile = None
+        self._recommendation = None
+
+    def profile(self):
+        """
+        Analyze the dataset.
+        """
+        if self._profile is None:
+            self._profile = DatasetProfiler(
+                self.df,
+                target=self.target
+            ).profile()
+
+        return self._profile
+
+    def recommend(self):
+        """
+        Recommend an appropriate splitting strategy.
+        """
+        if self._recommendation is None:
+            profile = self.profile()
+
+            self._recommendation = SplitRecommender().recommend(
+                profile
+            )
+
+        return self._recommendation
+
+    def split(
+        self,
+        recommendation=None,
+        test_size=0.2,
+        random_state=42,
+    ):
+        """
+        Execute the recommended train/test split.
+        """
+
+        if recommendation is None:
+            recommendation = self.recommend()
+
+        splitter = DatasetSplitter()
+
+        return splitter.split(
+            self.df,
+            target=self.target,
+            strategy=recommendation.strategy,
+            group_column=recommendation.group_column,
+            time_column=recommendation.time_column,
+            test_size=test_size,
+            random_state=random_state,
+        )

splitpilot/core/profiler.py

@@ -0,0 +1,207 @@
+from dataclasses import dataclass
+
+import pandas as pd
+
+
+@dataclass
+class GroupCandidate:
+    """A column that may represent a meaningful group or entity."""
+
+    column: str
+    score: float
+    reasons: list[str]
+
+
+@dataclass
+class DatasetProfile:
+    """Summary of the dataset relevant to split selection."""
+
+    rows: int
+    columns: int
+    target: str
+    target_type: str
+    target_unique_values: int
+    target_distribution: dict[str, float]
+    group_candidates: list[GroupCandidate]
+    possible_time_columns: list[str]
+    duplicate_rows: int
+
+
+class DatasetProfiler:
+    """Analyze a dataset before selecting an evaluation split."""
+
+    GROUP_NAME_HINTS = {
+        "id",
+        "customer",
+        "user",
+        "patient",
+        "account",
+        "member",
+        "client",
+        "session",
+        "device",
+        "store",
+        "merchant",
+        "household",
+        "group",
+        "entity",
+    }
+
+    def __init__(self, df: pd.DataFrame, target: str) -> None:
+        if not isinstance(df, pd.DataFrame):
+            raise TypeError("df must be a pandas DataFrame.")
+
+        if target not in df.columns:
+            raise ValueError(
+                f"Target column '{target}' was not found in the dataset."
+            )
+
+        if df.empty:
+            raise ValueError("The dataset cannot be empty.")
+
+        self.df = df
+        self.target = target
+
+    def profile(self) -> DatasetProfile:
+        """Create a structural profile of the dataset."""
+
+        target_series = self.df[self.target]
+
+        target_type = self._detect_target_type(target_series)
+
+        distribution = (
+            target_series.value_counts(normalize=True, dropna=False)
+            .round(4)
+            .to_dict()
+        )
+
+        distribution = {
+            str(key): float(value)
+            for key, value in distribution.items()
+        }
+
+        return DatasetProfile(
+            rows=len(self.df),
+            columns=len(self.df.columns),
+            target=self.target,
+            target_type=target_type,
+            target_unique_values=target_series.nunique(dropna=True),
+            target_distribution=distribution,
+            group_candidates=self._find_group_candidates(),
+            possible_time_columns=self._find_time_columns(),
+            duplicate_rows=int(self.df.duplicated().sum()),
+        )
+
+    def _detect_target_type(self, series: pd.Series) -> str:
+        """Determine whether the target looks like classification or regression."""
+
+        unique_values = series.nunique(dropna=True)
+
+        if pd.api.types.is_numeric_dtype(series):
+            if unique_values <= 20:
+                return "classification"
+
+            return "regression"
+
+        return "classification"
+
+    def _find_group_candidates(self) -> list[GroupCandidate]:
+        """Find and score columns that may represent groups or entities."""
+
+        candidates = []
+
+        for column in self.df.columns:
+            if column == self.target:
+                continue
+
+            series = self.df[column]
+
+            if series.isna().all():
+                continue
+
+            unique_count = series.nunique(dropna=True)
+
+            if unique_count <= 1:
+                continue
+
+            score = 0.0
+            reasons = []
+
+            # Signal 1: column name looks like an identifier/entity column.
+            name_parts = (
+                column.lower()
+                .replace("-", "_")
+                .split("_")
+            )
+
+            name_matches = self.GROUP_NAME_HINTS.intersection(name_parts)
+
+            if name_matches:
+                score += 0.60
+                reasons.append(
+                    "identifier/entity-like column name"
+                )
+
+            # Signal 2: values repeat across multiple rows.
+            repetition_ratio = 1 - (unique_count / len(self.df))
+
+            if repetition_ratio > 0:
+                score += min(repetition_ratio * 0.50, 0.25)
+                reasons.append("repeated values detected")
+
+            # Signal 3: each group has multiple observations.
+            group_sizes = series.value_counts(dropna=True)
+
+            if len(group_sizes) > 0:
+                average_group_size = group_sizes.mean()
+
+                if average_group_size >= 2:
+                    score += 0.15
+                    reasons.append("multiple observations per group")
+
+            score = min(score, 1.0)
+
+            # Only expose candidates with meaningful evidence.
+            if score >= 0.50:
+                candidates.append(
+                    GroupCandidate(
+                        column=column,
+                        score=round(score, 2),
+                        reasons=reasons,
+                    )
+                )
+
+        candidates.sort(
+            key=lambda candidate: candidate.score,
+            reverse=True,
+        )
+
+        return candidates
+
+    def _find_time_columns(self) -> list[str]:
+        """Find columns that appear to contain temporal information."""
+
+        candidates = []
+
+        for column in self.df.columns:
+            if column == self.target:
+                continue
+
+            series = self.df[column]
+
+            if pd.api.types.is_datetime64_any_dtype(series):
+                candidates.append(column)
+                continue
+
+            if pd.api.types.is_object_dtype(series):
+                parsed = pd.to_datetime(series, errors="coerce")
+
+                non_null = series.notna().sum()
+
+                if non_null > 0:
+                    parse_ratio = parsed.notna().sum() / non_null
+
+                    if parse_ratio >= 0.9:
+                        candidates.append(column)
+
+        return candidates

splitpilot/core/recommender.py

@@ -0,0 +1,127 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from .profiler import DatasetProfile
+
+
+@dataclass
+class SplitRecommendation:
+    strategy: str
+    reason: str
+    group_column: Optional[str] = None
+    time_column: Optional[str] = None
+
+
+class SplitRecommender:
+    """Recommend a train/test splitting strategy from a DatasetProfile."""
+
+    def recommend(self, profile: DatasetProfile) -> SplitRecommendation:
+
+        has_group = bool(profile.group_candidates)
+        has_time = bool(profile.possible_time_columns)
+
+        group_column = (
+            profile.group_candidates[0].column
+            if has_group
+            else None
+        )
+
+        time_column = (
+            profile.possible_time_columns[0]
+            if has_time
+            else None
+        )
+
+        # ---------------------------------------------------------
+        # 1. GROUP + TIME
+        # ---------------------------------------------------------
+        if has_group and has_time:
+
+            return SplitRecommendation(
+                strategy="group_time",
+                reason=(
+                    f"Repeated observations were detected for "
+                    f"'{group_column}', and temporal information was "
+                    f"detected in '{time_column}'. A split should respect "
+                    "both entity boundaries and chronological ordering "
+                    "to reduce leakage."
+                ),
+                group_column=group_column,
+                time_column=time_column,
+            )
+
+        # ---------------------------------------------------------
+        # 2. GROUP
+        # ---------------------------------------------------------
+        if has_group:
+
+            if profile.target_type == "classification":
+                strategy = "group_stratified"
+            else:
+                strategy = "group"
+
+            return SplitRecommendation(
+                strategy=strategy,
+                reason=(
+                    f"Multiple observations were detected for "
+                    f"'{group_column}'. A row-wise random split could "
+                    "place observations from the same entity in both "
+                    "train and test sets."
+                ),
+                group_column=group_column,
+            )
+
+        # ---------------------------------------------------------
+        # 3. TIME
+        # ---------------------------------------------------------
+        if has_time:
+
+            return SplitRecommendation(
+                strategy="time",
+                reason=(
+                    f"Temporal information was detected in "
+                    f"'{time_column}'. A chronological split can better "
+                    "represent the real-world situation where future "
+                    "observations are predicted from past observations."
+                ),
+                time_column=time_column,
+            )
+
+        # ---------------------------------------------------------
+        # 4. CLASSIFICATION
+        # ---------------------------------------------------------
+        if profile.target_type == "classification":
+
+            return SplitRecommendation(
+                strategy="stratified",
+                reason=(
+                    "The target is categorical and no group or temporal "
+                    "structure was detected. Stratification helps preserve "
+                    "the target distribution across train and test sets."
+                ),
+            )
+
+        # ---------------------------------------------------------
+        # 5. REGRESSION
+        # ---------------------------------------------------------
+        if profile.target_type == "regression":
+
+            return SplitRecommendation(
+                strategy="random",
+                reason=(
+                    "The target is continuous and no group or temporal "
+                    "structure was detected. A random split is therefore "
+                    "a reasonable default."
+                ),
+            )
+
+        # ---------------------------------------------------------
+        # 6. FALLBACK
+        # ---------------------------------------------------------
+        return SplitRecommendation(
+            strategy="random",
+            reason=(
+                "No specialized splitting structure was detected. "
+                "Falling back to a random split."
+            ),
+        )