Fix: handle single-class classifier in threshold prediction (#14471) (#14473)

scikit-learn treats a single-class fit as binary, so predict_proba returns
[1 - p, p] for one class and argmax could index past classes_. Use only
column 0, the real class, so the confidence threshold still applies.

src/documents/classifier.py

@@ -54,6 +54,10 @@ def _predict_with_threshold(classifier, X, threshold: float) -> int | None:
     cutoff so that uncertain predictions are discarded rather than assigned.
     """
     probas = classifier.predict_proba(X)[0]
+    if len(classifier.classes_) == 1:
+        # scikit-learn treats a single-class fit as binary, so predict_proba
+        # returns [1 - p, p] for one class. Column 0 is the real class.
+        probas = probas[:1]
     best_idx = int(probas.argmax())
     best_class = int(classifier.classes_[best_idx])
 

src/documents/tests/test_classifier.py

@@ -14,6 +14,7 @@ from django.test import override_settings
 from django.test.utils import CaptureQueriesContext
 from pytest_django.fixtures import Settings
 from pytest_mock import MockerFixture
+from sklearn.neural_network import MLPClassifier
 
 from documents.classifier import ClassifierModelCorruptError
 from documents.classifier import DocumentClassifier
@@ -1036,6 +1037,49 @@ def test_classifier_match_threshold_default() -> None:
     assert settings.CLASSIFIER_MATCH_THRESHOLD == 0.3
 
 
+class TestPredictWithThresholdSingleClass:
+    @pytest.fixture
+    def single_class_model(self) -> tuple[MLPClassifier, np.ndarray]:
+        X = np.random.default_rng(0).random((2, 5))
+        return MLPClassifier(tol=0.01, random_state=0).fit(X, [3, 3]), X[:1]
+
+    def test_returns_only_class(
+        self,
+        single_class_model: tuple[MLPClassifier, np.ndarray],
+    ) -> None:
+        """
+        GIVEN:
+            - A classifier fitted on a single class (every document shares one label)
+        WHEN:
+            - A prediction is requested with the threshold disabled
+        THEN:
+            - No error is raised and the only class is returned
+        """
+        classifier, X = single_class_model
+
+        assert _predict_with_threshold(classifier, X, threshold=0.0) == 3
+
+    def test_applies_threshold(
+        self,
+        single_class_model: tuple[MLPClassifier, np.ndarray],
+    ) -> None:
+        """
+        GIVEN:
+            - A classifier fitted on a single class
+        WHEN:
+            - A prediction is requested
+        THEN:
+            - The class's confidence is still compared against the threshold
+        """
+        classifier, X = single_class_model
+        confidence = classifier.predict_proba(X)[0][0]
+
+        assert _predict_with_threshold(classifier, X, threshold=confidence - 0.01) == 3
+        assert (
+            _predict_with_threshold(classifier, X, threshold=confidence + 0.01) is None
+        )
+
+
 class TestPreprocessContent:
     @pytest.fixture
     def samples(self) -> Path: