fix(models): validate inputs upfront and fix pickling of the Lightning module

deeptab/models/_mixins/fit.py

@@ -438,7 +438,11 @@ def fit(
             "fit.started",
             model_class=type(self).__name__,
             n_samples=len(X),
-            n_features=X.shape[1] if hasattr(X, "shape") else len(X.columns),
+            # X may still be a plain list here (ensure_dataframe runs in
+            # _build_model), so fall back to the first row's length.
+            n_features=(
+                X.shape[1] if hasattr(X, "shape") else (len(X.columns) if hasattr(X, "columns") else len(X[0]))
+            ),
             random_state=getattr(self, "random_state", None),
         )
 

deeptab/models/base.py

@@ -368,7 +368,7 @@ def __sklearn_is_fitted__(self) -> bool:
 
     def __getstate__(self):
         state = self.__dict__.copy()
-        state["task_model"] = None  # Avoid serializing the task model
+        state["_task_model"] = None  # Avoid serializing the Lightning module
         return state
 
     def __setstate__(self, state):

deeptab/models/classifier_base.py

@@ -28,6 +28,13 @@ def _resolve_loss_and_sampler(loss_fct, class_weight, balanced_sampler, sample_w
     resolved_loss = build_classification_loss(loss_fct, num_classes=num_classes, class_weights=class_weights)
 
     if sample_weight is not None:
+        weights = np.asarray(sample_weight, dtype=np.float64)
+        if (weights < 0).any():
+            raise ValueError("sample_weight must be non-negative.")
+        if not (weights > 0).any():
+            # Fail here with a clear message instead of a cryptic
+            # torch.multinomial RuntimeError mid-training.
+            raise ValueError("Sample weights must contain at least one non-zero value; all weights are zero.")
         sampler = sample_weight
     elif balanced_sampler:
         sampler = "balanced"

tests/test_api_robustness.py

@@ -0,0 +1,64 @@
+"""Regression tests for public-API robustness gaps found in the 2.0.0 review."""
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from deeptab.models import MLPClassifier, MLPRegressor
+
+FIT_KW = {"max_epochs": 1, "batch_size": 16, "accelerator": "cpu"}
+
+
+def _data(n=40, seed=0):
+    rng = np.random.RandomState(seed)
+    return pd.DataFrame({"a": rng.randn(n), "b": rng.randn(n)}), rng.randn(n)
+
+
+def test_getstate_excludes_lightning_module():
+    """__getstate__ nulled a phantom 'task_model' key, so pickling kept the real one."""
+    X, y = _data()
+    model = MLPRegressor()
+    model.fit(X, y, **FIT_KW)
+
+    state = model.__getstate__()
+    assert state["_task_model"] is None
+    assert state.get("task_model") is None
+
+
+def test_fit_accepts_plain_lists():
+    """The fit.started event read X.columns before ensure_dataframe ran."""
+    X = [[1.0, 2.0], [2.0, 1.0], [0.5, 0.1], [1.5, 0.7]] * 10
+    y = list(np.random.RandomState(0).randn(40))
+
+    model = MLPRegressor()
+    model.fit(X, y, **FIT_KW)
+    assert model.predict(X).shape == (40,)
+
+
+class TestSampleWeightValidation:
+    def test_all_zero_sample_weight_raises(self):
+        """Previously a cryptic torch.multinomial RuntimeError mid-training."""
+        X, _ = _data()
+        y = np.random.RandomState(1).choice([0, 1], len(X))
+
+        with pytest.raises(ValueError, match="zero"):
+            MLPClassifier().fit(X, y, sample_weight=np.zeros(len(X)), **FIT_KW)
+
+    def test_negative_sample_weight_raises(self):
+        X, _ = _data()
+        y = np.random.RandomState(1).choice([0, 1], len(X))
+        weights = np.ones(len(X))
+        weights[0] = -1.0
+
+        with pytest.raises(ValueError, match="non-negative"):
+            MLPClassifier().fit(X, y, sample_weight=weights, **FIT_KW)
+
+    def test_valid_sample_weight_still_accepted(self):
+        X, _ = _data()
+        y = np.random.RandomState(1).choice([0, 1], len(X))
+        weights = np.ones(len(X))
+        weights[:5] = 3.0
+
+        model = MLPClassifier()
+        model.fit(X, y, sample_weight=weights, **FIT_KW)
+        assert model.predict(X).shape == (len(X),)