diff --git a/src/documents/classifier.py b/src/documents/classifier.py index 3db7c3cfa..1cdef4419 100644 --- a/src/documents/classifier.py +++ b/src/documents/classifier.py @@ -19,6 +19,7 @@ if TYPE_CHECKING: from django.conf import settings from django.core.cache import cache from django.core.cache import caches +from django.db.models import Prefetch from documents.caching import CACHE_5_MINUTES from documents.caching import CACHE_50_MINUTES @@ -28,6 +29,7 @@ from documents.caching import CLASSIFIER_VERSION_KEY from documents.caching import StoredLRUCache from documents.models import Document from documents.models import MatchingModel +from documents.models import Tag from paperless.signed_pickle import SignedPickleError from paperless.signed_pickle import signed_pickle_dumps from paperless.signed_pickle import signed_pickle_loads @@ -260,7 +262,17 @@ class DocumentClassifier: tags__is_inbox_tag=True, ) .select_related("document_type", "correspondent", "storage_path") - .prefetch_related("tags") + .prefetch_related( + Prefetch( + "tags", + queryset=Tag.objects.filter( + matching_algorithm=MatchingModel.MATCH_AUTO, + ) + .order_by("pk") + .only("pk"), + to_attr="auto_tags", + ), + ) .order_by("pk") ) @@ -292,11 +304,7 @@ class DocumentClassifier: hasher.update(y.to_bytes(4, "little", signed=True)) labels_correspondent.append(y) - tags: list[int] = list( - doc.tags.filter(matching_algorithm=MatchingModel.MATCH_AUTO) - .order_by("pk") - .values_list("pk", flat=True), - ) + tags: list[int] = [tag.pk for tag in doc.auto_tags] for tag in tags: hasher.update(tag.to_bytes(4, "little", signed=True)) labels_tags.append(tags) diff --git a/src/documents/tests/test_classifier.py b/src/documents/tests/test_classifier.py index 0e09ff074..9066c4d41 100644 --- a/src/documents/tests/test_classifier.py +++ b/src/documents/tests/test_classifier.py @@ -6,8 +6,11 @@ from unittest import mock import numpy as np import pytest from django.conf import settings +from django.db import connection from django.test import TestCase from django.test import override_settings +from django.test.utils import CaptureQueriesContext +from pytest_mock import MockerFixture from documents.classifier import ClassifierModelCorruptError from documents.classifier import DocumentClassifier @@ -20,6 +23,8 @@ from documents.models import DocumentType from documents.models import MatchingModel from documents.models import StoragePath from documents.models import Tag +from documents.tests.factories import DocumentFactory +from documents.tests.factories import TagFactory from documents.tests.utils import DirectoriesMixin from paperless.signed_pickle import HMAC_SIZE from paperless.signed_pickle import signed_pickle_dumps @@ -1017,3 +1022,70 @@ def test_preprocess_content_nltk_load_fail(mocker) -> None: expected_preprocess_content = f.read().rstrip() result = classifier.preprocess_content(content) assert result == expected_preprocess_content + + +@pytest.mark.django_db +class TestClassifierTrainTagLabels: + @pytest.fixture(autouse=True) + def _simple_preprocess(self, mocker: MockerFixture) -> None: + mocker.patch.object( + DocumentClassifier, + "preprocess_content", + side_effect=dummy_preprocess, + ) + + @pytest.fixture + def auto_tags(self) -> list[Tag]: + return TagFactory.create_batch(2, matching_algorithm=MatchingModel.MATCH_AUTO) + + def test_train_query_count_does_not_scale_with_documents( + self, + auto_tags: list[Tag], + ) -> None: + """ + GIVEN: + - Documents with auto matching tags + WHEN: + - The classifier is trained, then more documents are added and it is + trained again + THEN: + - Both trainings run the same number of queries + """ + for doc in DocumentFactory.create_batch(2): + doc.tags.set(auto_tags) + + with CaptureQueriesContext(connection) as few_documents: + DocumentClassifier().train() + + for doc in DocumentFactory.create_batch(6): + doc.tags.set(auto_tags) + + with CaptureQueriesContext(connection) as more_documents: + DocumentClassifier().train() + + assert len(more_documents) == len(few_documents) + + def test_train_uses_only_auto_tags_as_labels( + self, + auto_tags: list[Tag], + ) -> None: + """ + GIVEN: + - Documents with both auto matching and non auto matching tags + WHEN: + - The classifier is trained + THEN: + - Only the auto matching tags are used as tag labels + """ + manual_tag = TagFactory(matching_algorithm=MatchingModel.MATCH_ANY) + first, second, third = DocumentFactory.create_batch(3) + first.tags.set([auto_tags[0], manual_tag]) + second.tags.set([auto_tags[1], manual_tag]) + third.tags.set([*auto_tags, manual_tag]) + + classifier = DocumentClassifier() + classifier.train() + + assert list(classifier.tags_binarizer.classes_) == sorted( + tag.pk for tag in auto_tags + )