Enhancement: Improve matching for correspondents, storage path and labels by removing bias + adding minimum match threshold (#12164)

docs/configuration.md

@@ -1200,6 +1200,15 @@ still perform some basic text pre-processing before matching.
 
     Defaults to true, enabling the feature.
 
+#### [`PAPERLESS_CLASSIFIER_MATCH_THRESHOLD=<float>`](#PAPERLESS_CLASSIFIER_MATCH_THRESHOLD) {#PAPERLESS_CLASSIFIER_MATCH_THRESHOLD}
+
+: Sets the minimum confidence score (0.0-1.0) required for the automatic
+classifier to assign a correspondent, document type, or storage path to a
+document. Predictions below this threshold are discarded and the field is
+left unassigned, preventing low-confidence guesses from being applied.
+
+    Defaults to 0.6.
+
 #### [`PAPERLESS_DATE_PARSER_LANGUAGES=<lang>`](#PAPERLESS_DATE_PARSER_LANGUAGES) {#PAPERLESS_DATE_PARSER_LANGUAGES}
 
 : Specifies which language Paperless should use when parsing dates from documents.

src/documents/classifier.py

@@ -34,6 +34,27 @@ from paperless.signed_pickle import signed_pickle_loads
 
 logger = logging.getLogger("paperless.classifier")
 
+
+def _predict_with_threshold(classifier, X, threshold: float) -> int | None:
+    """
+    Return the predicted class id, or None if:
+    - the prediction is -1 (no match), or
+    - the winning class probability is below the configured threshold.
+
+    Using predict_proba() instead of predict() lets us apply a minimum-confidence
+    cutoff so that uncertain predictions are discarded rather than assigned.
+    """
+    probas = classifier.predict_proba(X)[0]
+    best_idx = int(probas.argmax())
+    best_class = int(classifier.classes_[best_idx])
+
+    if best_class == -1:
+        return None
+    if threshold > 0.0 and probas[best_idx] < threshold:
+        return None
+    return best_class
+
+
 ADVANCED_TEXT_PROCESSING_ENABLED = (
     settings.NLTK_LANGUAGE is not None and settings.NLTK_ENABLED
 )
@@ -102,7 +123,8 @@ class DocumentClassifier:
     # v8 - Added storage path classifier
     # v9 - Changed from hashing to time/ids for re-train check
     # v10 - HMAC-signed model file
-    FORMAT_VERSION = 10
+    # v11 - Use sample_weight for balanced training; predict_proba with threshold
+    FORMAT_VERSION = 11
 
     HMAC_SIZE = 32  # SHA-256 digest length
 
@@ -324,6 +346,13 @@ class DocumentClassifier:
         from sklearn.preprocessing import LabelBinarizer
         from sklearn.preprocessing import MultiLabelBinarizer
 
+        # MLPClassifier does not support class_weight directly
+        # (https://github.com/scikit-learn/scikit-learn/issues/9113), so we use
+        # compute_sample_weight to balance classes during training and prevent
+        # over-represented correspondents from dominating predictions.
+        # https://scikit-learn.org/stable/modules/generated/sklearn.utils.class_weight.compute_sample_weight.html
+        from sklearn.utils.class_weight import compute_sample_weight
+
         # Step 2: vectorize data
         logger.debug("Vectorizing data...")
         notify("Vectorizing document content...")
@@ -369,7 +398,7 @@ class DocumentClassifier:
                 self.tags_binarizer = MultiLabelBinarizer()
                 labels_tags_vectorized = self.tags_binarizer.fit_transform(labels_tags)
 
-            self.tags_classifier = MLPClassifier(tol=0.01)
+            self.tags_classifier = MLPClassifier(tol=0.01, random_state=0)
             self.tags_classifier.fit(data_vectorized, labels_tags_vectorized)
         else:
             self.tags_classifier = None
@@ -380,8 +409,12 @@ class DocumentClassifier:
             notify(
                 f"Training correspondent classifier ({num_correspondents} correspondent(s))...",
             )
-            self.correspondent_classifier = MLPClassifier(tol=0.01)
-            self.correspondent_classifier.fit(data_vectorized, labels_correspondent)
+            self.correspondent_classifier = MLPClassifier(tol=0.01, random_state=0)
+            self.correspondent_classifier.fit(
+                data_vectorized,
+                labels_correspondent,
+                sample_weight=compute_sample_weight("balanced", labels_correspondent),
+            )
         else:
             self.correspondent_classifier = None
             logger.debug(
@@ -393,8 +426,12 @@ class DocumentClassifier:
             notify(
                 f"Training document type classifier ({num_document_types} type(s))...",
             )
-            self.document_type_classifier = MLPClassifier(tol=0.01)
-            self.document_type_classifier.fit(data_vectorized, labels_document_type)
+            self.document_type_classifier = MLPClassifier(tol=0.01, random_state=0)
+            self.document_type_classifier.fit(
+                data_vectorized,
+                labels_document_type,
+                sample_weight=compute_sample_weight("balanced", labels_document_type),
+            )
         else:
             self.document_type_classifier = None
             logger.debug(
@@ -406,10 +443,11 @@ class DocumentClassifier:
                 "Training storage paths classifier...",
             )
             notify(f"Training storage path classifier ({num_storage_paths} path(s))...")
-            self.storage_path_classifier = MLPClassifier(tol=0.01)
+            self.storage_path_classifier = MLPClassifier(tol=0.01, random_state=0)
             self.storage_path_classifier.fit(
                 data_vectorized,
                 labels_storage_path,
+                sample_weight=compute_sample_weight("balanced", labels_storage_path),
             )
         else:
             self.storage_path_classifier = None
@@ -546,24 +584,24 @@ class DocumentClassifier:
     def predict_correspondent(self, content: str) -> int | None:
         if self.correspondent_classifier:
             X = self._vectorize(content)
-            correspondent_id = self.correspondent_classifier.predict(X)
-            if correspondent_id != -1:
-                return correspondent_id
-            else:
-                return None
-        else:
-            return None
+            predicted_id = _predict_with_threshold(
+                self.correspondent_classifier,
+                X,
+                settings.CLASSIFIER_MATCH_THRESHOLD,
+            )
+            return predicted_id
+        return None
 
     def predict_document_type(self, content: str) -> int | None:
         if self.document_type_classifier:
             X = self._vectorize(content)
-            document_type_id = self.document_type_classifier.predict(X)
-            if document_type_id != -1:
-                return document_type_id
-            else:
-                return None
-        else:
-            return None
+            predicted_id = _predict_with_threshold(
+                self.document_type_classifier,
+                X,
+                settings.CLASSIFIER_MATCH_THRESHOLD,
+            )
+            return predicted_id
+        return None
 
     def predict_tags(self, content: str) -> list[int]:
         from sklearn.utils.multiclass import type_of_target
@@ -589,10 +627,10 @@ class DocumentClassifier:
     def predict_storage_path(self, content: str) -> int | None:
         if self.storage_path_classifier:
             X = self._vectorize(content)
-            storage_path_id = self.storage_path_classifier.predict(X)
-            if storage_path_id != -1:
-                return storage_path_id
-            else:
-                return None
-        else:
-            return None
+            predicted_id = _predict_with_threshold(
+                self.storage_path_classifier,
+                X,
+                settings.CLASSIFIER_MATCH_THRESHOLD,
+            )
+            return predicted_id
+        return None

src/documents/tests/test_classifier.py

@@ -3,6 +3,7 @@ import warnings
 from pathlib import Path
 from unittest import mock
 
+import numpy as np
 import pytest
 from django.conf import settings
 from django.test import TestCase
@@ -11,6 +12,7 @@ from django.test import override_settings
 from documents.classifier import ClassifierModelCorruptError
 from documents.classifier import DocumentClassifier
 from documents.classifier import IncompatibleClassifierVersionError
+from documents.classifier import _predict_with_threshold
 from documents.classifier import load_classifier
 from documents.models import Correspondent
 from documents.models import Document
@@ -625,6 +627,103 @@ class TestClassifier(DirectoriesMixin, TestCase):
         self.assertEqual(self.classifier.predict_storage_path(doc1.content), sp.pk)
         self.assertIsNone(self.classifier.predict_storage_path(doc2.content))
 
+    def test_predict_rejects_prediction_below_match_threshold(self) -> None:
+        """
+        GIVEN:
+            - Classifiers trained against test data with confident predictions
+        WHEN:
+            - CLASSIFIER_MATCH_THRESHOLD exceeds the model's confidence
+        THEN:
+            - Every predict_* method discards the match in favor of no match
+        """
+        c1 = Correspondent.objects.create(
+            name="c1",
+            matching_algorithm=Correspondent.MATCH_AUTO,
+        )
+        dt1 = DocumentType.objects.create(
+            name="dt1",
+            matching_algorithm=DocumentType.MATCH_AUTO,
+        )
+        sp1 = StoragePath.objects.create(
+            name="sp1",
+            matching_algorithm=StoragePath.MATCH_AUTO,
+        )
+
+        doc1 = Document.objects.create(
+            title="doc1",
+            content="this is a document from c1",
+            correspondent=c1,
+            document_type=dt1,
+            storage_path=sp1,
+            checksum="A",
+        )
+        Document.objects.create(
+            title="doc2",
+            content="this is a document from no one",
+            checksum="B",
+        )
+
+        self.classifier.train()
+
+        predictors = {
+            "correspondent": self.classifier.predict_correspondent,
+            "document_type": self.classifier.predict_document_type,
+            "storage_path": self.classifier.predict_storage_path,
+        }
+        # No real prediction can reach a confidence this high, so this
+        # isolates the threshold check from the model's actual output.
+        with override_settings(CLASSIFIER_MATCH_THRESHOLD=0.999999):
+            for name, predict in predictors.items():
+                with self.subTest(field=name):
+                    self.assertIsNone(predict(doc1.content))
+
+    def test_train_uses_balanced_sample_weight(self) -> None:
+        """
+        GIVEN:
+            - A training set with correspondents, document types and storage paths
+        WHEN:
+            - The classifier is trained
+        THEN:
+            - Each MLP classifier is fit with balanced sample weights, so that
+              over-represented classes don't dominate predictions
+        """
+        c1 = Correspondent.objects.create(
+            name="c1",
+            matching_algorithm=Correspondent.MATCH_AUTO,
+        )
+        dt1 = DocumentType.objects.create(
+            name="dt1",
+            matching_algorithm=DocumentType.MATCH_AUTO,
+        )
+        sp1 = StoragePath.objects.create(
+            name="sp1",
+            matching_algorithm=StoragePath.MATCH_AUTO,
+        )
+
+        Document.objects.create(
+            title="doc1",
+            content="this is a document from c1",
+            correspondent=c1,
+            document_type=dt1,
+            storage_path=sp1,
+            checksum="A",
+        )
+        Document.objects.create(
+            title="doc2",
+            content="this is a document from no one",
+            checksum="B",
+        )
+
+        with mock.patch(
+            "sklearn.utils.class_weight.compute_sample_weight",
+            return_value=None,
+        ) as mocked_compute_sample_weight:
+            self.classifier.train()
+
+        self.assertEqual(mocked_compute_sample_weight.call_count, 3)
+        for call in mocked_compute_sample_weight.call_args_list:
+            self.assertEqual(call.args[0], "balanced")
+
     def test_one_tag_predict(self) -> None:
         t1 = Tag.objects.create(name="t1", matching_algorithm=Tag.MATCH_AUTO, pk=12)
 
@@ -810,6 +909,52 @@ class TestClassifier(DirectoriesMixin, TestCase):
             load_classifier(raise_exception=True)
 
 
+class _StubProbaClassifier:
+    """
+    A fake scikit-learn classifier exposing just enough of the API for
+    `_predict_with_threshold`: `classes_` and `predict_proba`.
+    """
+
+    def __init__(self, classes: list[int], probabilities: list[float]) -> None:
+        self.classes_ = np.array(classes)
+        self._probabilities = np.array([probabilities])
+
+    def predict_proba(self, X) -> np.ndarray:
+        return self._probabilities
+
+
+@pytest.mark.parametrize(
+    ("classes", "probabilities", "threshold", "expected"),
+    [
+        # confident prediction above the threshold is returned
+        ([-1, 3], [0.1, 0.9], 0.6, 3),
+        # prediction below the threshold is discarded
+        ([-1, 3], [0.45, 0.55], 0.6, None),
+        # boundary: exactly at the threshold is accepted, not discarded
+        ([-1, 3], [0.4, 0.6], 0.6, 3),
+        # the winning class is the "no match" pseudo-class, regardless of its
+        # own confidence
+        ([-1, 3], [0.99, 0.01], 0.0, None),
+        # threshold of 0.0 disables the confidence check entirely
+        ([-1, 3], [0.45, 0.55], 0.0, 3),
+    ],
+)
+def test_predict_with_threshold(classes, probabilities, threshold, expected) -> None:
+    classifier = _StubProbaClassifier(classes, probabilities)
+    result = _predict_with_threshold(classifier, X=None, threshold=threshold)
+    assert result == expected
+
+
+def test_classifier_match_threshold_default() -> None:
+    """
+    GIVEN:
+        - No PAPERLESS_CLASSIFIER_MATCH_THRESHOLD environment variable is set
+    THEN:
+        - The classifier match threshold defaults to 0.6
+    """
+    assert settings.CLASSIFIER_MATCH_THRESHOLD == 0.6
+
+
 def test_preprocess_content() -> None:
     """
     GIVEN:

src/paperless/settings/__init__.py

@@ -96,6 +96,13 @@ MODEL_FILE = get_path_from_env(
     "PAPERLESS_MODEL_FILE",
     DATA_DIR / "classification_model.pickle",
 )
+
+# Minimum confidence (0.0-1.0) for the ML classifier to assign a correspondent,
+# document type, or storage path. 0.0 disables the threshold.
+CLASSIFIER_MATCH_THRESHOLD: Final[float] = get_float_from_env(
+    "PAPERLESS_CLASSIFIER_MATCH_THRESHOLD",
+    0.6,
+)
 LLM_INDEX_DIR = DATA_DIR / "llm_index"
 LLM_INDEX_LOCK = LLM_INDEX_DIR / "index.lock"
 # Cross-process read/write lock guarding the LLM index compaction/migration