From 0cb2cdbfdd91a2078236f8851cd310261284591c Mon Sep 17 00:00:00 2001
From: Trenton H <797416+stumpylog@users.noreply.github.com>
Date: Fri, 16 Oct 2026 10:47:58 -0700
Subject: [PATCH] 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            |  4 +++
 src/documents/tests/test_classifier.py | 44 ++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/src/documents/classifier.py b/src/documents/classifier.py
index e98b8b574..9c5ecaa62 100644
--- a/src/documents/classifier.py
+++ b/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])
 
diff --git a/src/documents/tests/test_classifier.py b/src/documents/tests/test_classifier.py
index ae8921ad6..c8b832a6a 100644
--- a/src/documents/tests/test_classifier.py
+++ b/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: