diff --git a/docs/configuration.md b/docs/configuration.md index 37a552f1a..e045af077 100644 --- a/docs/configuration.md +++ b/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=`](#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=`](#PAPERLESS_DATE_PARSER_LANGUAGES) {#PAPERLESS_DATE_PARSER_LANGUAGES} : Specifies which language Paperless should use when parsing dates from documents. diff --git a/src/documents/classifier.py b/src/documents/classifier.py index a5c272062..cfed4725c 100644 --- a/src/documents/classifier.py +++ b/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 diff --git a/src/documents/tests/test_classifier.py b/src/documents/tests/test_classifier.py index 12dfea4d8..0e09ff074 100644 --- a/src/documents/tests/test_classifier.py +++ b/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: diff --git a/src/paperless/settings/__init__.py b/src/paperless/settings/__init__.py index a0dc0e585..14f2fb9ac 100644 --- a/src/paperless/settings/__init__.py +++ b/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