mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-10-17 13:28:12 +00:00
Trenton H
530059c5c0
* Preprocesses classifier content with Tantivy instead of NLTK Tokenizing and stemming now happen in one Rust call instead of NLTK's Python tokenizer and per word stemming, which also removes the Redis backed stem cache from every preprocessing call. The output matches the NLTK pipeline closely; tokens containing digits are now stemmed, and the English stop words follow Snowball's list. Stemming and stop word removal apply whenever the OCR language is one of the supported classifier languages, so PAPERLESS_ENABLE_NLTK and PAPERLESS_NLTK_DIR are removed. * Copies packages instead of hardlinking them in backend CI, some NLTK thing * Adds a normalization to NFC to better fit what Tantivy expects
673 lines
26 KiB
Python
673 lines
26 KiB
Python
from __future__ import annotations
|
|
|
|
import functools
|
|
import hmac
|
|
import logging
|
|
import pickle
|
|
import re
|
|
import unicodedata
|
|
import warnings
|
|
from hashlib import sha256
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
from collections.abc import Iterator
|
|
from datetime import datetime
|
|
from types import TracebackType
|
|
from typing import BinaryIO
|
|
from typing import Self
|
|
|
|
import tantivy
|
|
from numpy import ndarray
|
|
from sklearn.neural_network import MLPClassifier
|
|
|
|
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._snowball_stopwords import ENGLISH as ENGLISH_STOP_WORDS
|
|
from documents.caching import CACHE_5_MINUTES
|
|
from documents.caching import CACHE_50_MINUTES
|
|
from documents.caching import CLASSIFIER_HASH_KEY
|
|
from documents.caching import CLASSIFIER_MODIFIED_KEY
|
|
from documents.caching import CLASSIFIER_VERSION_KEY
|
|
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
|
|
|
|
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
|
|
|
|
|
|
read_cache = caches["read-cache"]
|
|
|
|
|
|
RE_WORD = re.compile(r"\b[\w]+\b") # words that may contain digits
|
|
|
|
# Documents whose content is fetched per query while training
|
|
_CONTENT_CHUNK_SIZE = 1000
|
|
|
|
|
|
class _SignedFileWriter:
|
|
"""
|
|
Atomically writes a file made of an HMAC signature followed by the data,
|
|
signing the data as it streams to disk rather than holding it in memory.
|
|
|
|
The signature is only known once everything is written, so its space is
|
|
reserved at the start of the file and filled in on exit. The target is only
|
|
replaced once the file is complete; on error the partial file is removed.
|
|
"""
|
|
|
|
def __init__(self, target: Path, mac: hmac.HMAC) -> None:
|
|
self._target = target
|
|
self._temp = target.with_name(f"{target.name}.part")
|
|
self._mac = mac
|
|
self._file: BinaryIO
|
|
|
|
def __enter__(self) -> Self:
|
|
self._file = self._temp.open("wb")
|
|
self._file.write(bytes(self._mac.digest_size))
|
|
return self
|
|
|
|
def write(self, data: bytes | memoryview) -> int:
|
|
self._mac.update(data)
|
|
return self._file.write(data)
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc_value: BaseException | None,
|
|
traceback: TracebackType | None,
|
|
) -> None:
|
|
try:
|
|
with self._file:
|
|
if exc_type is None:
|
|
self._file.seek(0)
|
|
self._file.write(self._mac.digest())
|
|
if exc_type is None:
|
|
self._temp.rename(self._target)
|
|
finally:
|
|
# A no-op after a successful rename, otherwise removes the partial file
|
|
self._temp.unlink(missing_ok=True)
|
|
|
|
|
|
@functools.cache
|
|
def _text_analyzer(language: str) -> tantivy.TextAnalyzer:
|
|
"""
|
|
Builds the cached analyzer for a language: word tokens, lowercase, stop words, stemmer.
|
|
|
|
Long tokens are kept and accents are not folded to ASCII, since stemmers
|
|
for languages such as French and German rely on accents.
|
|
"""
|
|
import tantivy
|
|
|
|
if language == "english":
|
|
# Tantivy's builtin English list is much shorter than Snowball's.
|
|
# Split contractions ("don't") on word characters, as content is, so they match
|
|
stop_words = tantivy.Filter.custom_stopword(
|
|
sorted({t for word in ENGLISH_STOP_WORDS for t in RE_WORD.findall(word)}),
|
|
)
|
|
else:
|
|
stop_words = tantivy.Filter.stopword(language)
|
|
return (
|
|
tantivy.TextAnalyzerBuilder(tantivy.Tokenizer.regex(r"\w+"))
|
|
.filter(tantivy.Filter.lowercase())
|
|
.filter(stop_words)
|
|
.filter(tantivy.Filter.stemmer(language))
|
|
.build()
|
|
)
|
|
|
|
|
|
class IncompatibleClassifierVersionError(Exception):
|
|
def __init__(self, message: str, *args: object) -> None:
|
|
self.message: str = message
|
|
super().__init__(*args)
|
|
|
|
|
|
class ClassifierModelCorruptError(Exception):
|
|
pass
|
|
|
|
|
|
def load_classifier(*, raise_exception: bool = False) -> DocumentClassifier | None:
|
|
if not settings.MODEL_FILE.is_file():
|
|
logger.debug(
|
|
"Document classification model does not exist (yet), not "
|
|
"performing automatic matching.",
|
|
)
|
|
return None
|
|
|
|
classifier = DocumentClassifier()
|
|
try:
|
|
classifier.load()
|
|
|
|
except IncompatibleClassifierVersionError as e:
|
|
logger.info(f"Classifier version incompatible: {e.message}, will re-train")
|
|
Path(settings.MODEL_FILE).unlink()
|
|
classifier = None
|
|
if raise_exception:
|
|
raise e
|
|
except ClassifierModelCorruptError as e:
|
|
# there's something wrong with the model file.
|
|
logger.exception(
|
|
"Unrecoverable error while loading document "
|
|
"classification model, deleting model file.",
|
|
)
|
|
Path(settings.MODEL_FILE).unlink()
|
|
classifier = None
|
|
if raise_exception:
|
|
raise e
|
|
except OSError as e:
|
|
logger.exception("IO error while loading document classification model")
|
|
classifier = None
|
|
if raise_exception:
|
|
raise e
|
|
except Exception as e: # pragma: no cover
|
|
logger.exception("Unknown error while loading document classification model")
|
|
classifier = None
|
|
if raise_exception:
|
|
raise e
|
|
|
|
return classifier
|
|
|
|
|
|
class DocumentClassifier:
|
|
# v7 - Updated scikit-learn package version
|
|
# v8 - Added storage path classifier
|
|
# v9 - Changed from hashing to time/ids for re-train check
|
|
# v10 - HMAC-signed model file
|
|
# v11 - Use sample_weight for balanced training; predict_proba with threshold;
|
|
# drop training-only MLP state before saving
|
|
# Tantivy text preprocessing
|
|
FORMAT_VERSION = 11
|
|
|
|
HMAC_SIZE = 32 # SHA-256 digest length
|
|
|
|
def __init__(self) -> None:
|
|
# last time a document changed and therefore training might be required
|
|
self.last_doc_change_time: datetime | None = None
|
|
# Hash of primary keys of AUTO matching values last used in training
|
|
self.last_auto_type_hash: bytes | None = None
|
|
|
|
self.data_vectorizer = None
|
|
self.data_vectorizer_hash = None
|
|
self.tags_binarizer = None
|
|
self.tags_classifier = None
|
|
self.correspondent_classifier = None
|
|
self.document_type_classifier = None
|
|
self.storage_path_classifier = None
|
|
|
|
def _update_data_vectorizer_hash(self) -> None:
|
|
self.data_vectorizer_hash = sha256(
|
|
pickle.dumps(self.data_vectorizer),
|
|
).hexdigest()
|
|
|
|
@staticmethod
|
|
def _new_hmac() -> hmac.HMAC:
|
|
return hmac.new(settings.SECRET_KEY.encode(), digestmod=sha256)
|
|
|
|
@staticmethod
|
|
def _strip_training_state(classifier: MLPClassifier) -> None:
|
|
"""
|
|
Drop MLPClassifier state which is only used during fit(), never by predict().
|
|
|
|
The Adam optimizer keeps two moment arrays the size of the weights, and
|
|
without early_stopping _best_coefs/_best_intercepts are just a copy of the
|
|
initial random weights. Together that is 3x the size of the weights, which
|
|
would otherwise be pickled and loaded along with the model.
|
|
"""
|
|
del classifier._optimizer
|
|
del classifier._best_coefs
|
|
del classifier._best_intercepts
|
|
|
|
@staticmethod
|
|
def _compute_hmac(data: bytes | memoryview) -> bytes:
|
|
mac = DocumentClassifier._new_hmac()
|
|
mac.update(data)
|
|
return mac.digest()
|
|
|
|
def load(self) -> None:
|
|
from sklearn.exceptions import InconsistentVersionWarning
|
|
|
|
raw = Path(settings.MODEL_FILE).read_bytes()
|
|
|
|
if len(raw) <= self.HMAC_SIZE:
|
|
raise ClassifierModelCorruptError
|
|
|
|
# Slice through a memoryview so the (potentially multi-GB) payload is
|
|
# not copied; hmac and pickle both accept buffers directly.
|
|
# The whole file is still verified from memory before unpickling, rather
|
|
# than streamed from disk, so it cannot change between check and load.
|
|
view = memoryview(raw)
|
|
signature = view[: self.HMAC_SIZE]
|
|
data = view[self.HMAC_SIZE :]
|
|
|
|
if not hmac.compare_digest(signature, self._compute_hmac(data)):
|
|
raise ClassifierModelCorruptError
|
|
|
|
# Catch warnings for processing
|
|
with warnings.catch_warnings(record=True) as w:
|
|
try:
|
|
(
|
|
schema_version,
|
|
self.last_doc_change_time,
|
|
self.last_auto_type_hash,
|
|
self.data_vectorizer,
|
|
self.tags_binarizer,
|
|
self.tags_classifier,
|
|
self.correspondent_classifier,
|
|
self.document_type_classifier,
|
|
self.storage_path_classifier,
|
|
) = pickle.loads(data)
|
|
except Exception as err:
|
|
raise ClassifierModelCorruptError from err
|
|
|
|
if schema_version != self.FORMAT_VERSION:
|
|
raise IncompatibleClassifierVersionError(
|
|
"Cannot load classifier, incompatible versions.",
|
|
)
|
|
|
|
self._update_data_vectorizer_hash()
|
|
|
|
# Check for the warning about unpickling from differing versions
|
|
# and consider it incompatible
|
|
sk_learn_warning_url = (
|
|
"https://scikit-learn.org/stable/"
|
|
"model_persistence.html"
|
|
"#security-maintainability-limitations"
|
|
)
|
|
for warning in w:
|
|
# The warning is inconsistent, the MLPClassifier is a specific warning, others have not updated yet
|
|
if issubclass(warning.category, InconsistentVersionWarning) or (
|
|
issubclass(warning.category, UserWarning)
|
|
and sk_learn_warning_url in str(warning.message)
|
|
):
|
|
raise IncompatibleClassifierVersionError("sklearn version update")
|
|
|
|
def save(self) -> None:
|
|
# Stream to disk instead of building the payload in memory. Protocol 5+
|
|
# pickles numpy arrays without copying them (the default is 4 before 3.14).
|
|
with _SignedFileWriter(settings.MODEL_FILE, self._new_hmac()) as f:
|
|
pickle.dump(
|
|
(
|
|
self.FORMAT_VERSION,
|
|
self.last_doc_change_time,
|
|
self.last_auto_type_hash,
|
|
self.data_vectorizer,
|
|
self.tags_binarizer,
|
|
self.tags_classifier,
|
|
self.correspondent_classifier,
|
|
self.document_type_classifier,
|
|
self.storage_path_classifier,
|
|
),
|
|
f,
|
|
protocol=pickle.HIGHEST_PROTOCOL,
|
|
)
|
|
|
|
def train(
|
|
self,
|
|
status_callback: Callable[[str], None] | None = None,
|
|
) -> bool:
|
|
notify = status_callback if status_callback is not None else lambda _: None
|
|
|
|
# Get non-inbox documents
|
|
docs_queryset = Document.objects.exclude(
|
|
tags__is_inbox_tag=True,
|
|
).order_by("pk")
|
|
|
|
# No documents exit to train against
|
|
doc_count = docs_queryset.count()
|
|
if doc_count == 0:
|
|
raise ValueError("No training data available.")
|
|
|
|
labels_tags = []
|
|
labels_correspondent = []
|
|
labels_document_type = []
|
|
labels_storage_path = []
|
|
# Content is fetched separately later, for exactly these documents in this
|
|
# order, so it never all has to be in memory at once
|
|
doc_pks: list[int] = []
|
|
latest_doc_change: datetime | None = None
|
|
|
|
# Step 1: Extract and preprocess training data from the database.
|
|
logger.debug("Gathering data from database...")
|
|
notify(f"Gathering data from {doc_count} document(s)...")
|
|
hasher = sha256()
|
|
for doc in (
|
|
docs_queryset.defer("content")
|
|
.select_related("document_type", "correspondent", "storage_path")
|
|
.prefetch_related(
|
|
Prefetch(
|
|
"tags",
|
|
queryset=Tag.objects.filter(
|
|
matching_algorithm=MatchingModel.MATCH_AUTO,
|
|
)
|
|
.order_by("pk")
|
|
.only("pk"),
|
|
to_attr="auto_tags",
|
|
),
|
|
)
|
|
.iterator(chunk_size=2000)
|
|
):
|
|
doc_pks.append(doc.pk)
|
|
if latest_doc_change is None or doc.modified > latest_doc_change:
|
|
latest_doc_change = doc.modified
|
|
|
|
y = -1
|
|
dt = doc.document_type
|
|
if dt and dt.matching_algorithm == MatchingModel.MATCH_AUTO:
|
|
y = dt.pk
|
|
hasher.update(y.to_bytes(4, "little", signed=True))
|
|
labels_document_type.append(y)
|
|
|
|
y = -1
|
|
cor = doc.correspondent
|
|
if cor and cor.matching_algorithm == MatchingModel.MATCH_AUTO:
|
|
y = cor.pk
|
|
hasher.update(y.to_bytes(4, "little", signed=True))
|
|
labels_correspondent.append(y)
|
|
|
|
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)
|
|
|
|
y = -1
|
|
sp = doc.storage_path
|
|
if sp and sp.matching_algorithm == MatchingModel.MATCH_AUTO:
|
|
y = sp.pk
|
|
hasher.update(y.to_bytes(4, "little", signed=True))
|
|
labels_storage_path.append(y)
|
|
|
|
labels_tags_unique = {tag for tags in labels_tags for tag in tags}
|
|
|
|
num_tags = len(labels_tags_unique)
|
|
|
|
# Check if retraining is actually required.
|
|
# A document has been updated since the classifier was trained
|
|
# New auto tags, types, correspondent, storage paths exist
|
|
if (
|
|
self.last_doc_change_time is not None
|
|
and self.last_doc_change_time >= latest_doc_change
|
|
) and self.last_auto_type_hash == hasher.digest():
|
|
logger.info("No updates since last training")
|
|
# Set the classifier information into the cache
|
|
# Caching for 50 minutes, so slightly less than the normal retrain time
|
|
cache.set(
|
|
CLASSIFIER_MODIFIED_KEY,
|
|
self.last_doc_change_time,
|
|
CACHE_50_MINUTES,
|
|
)
|
|
cache.set(CLASSIFIER_HASH_KEY, hasher.hexdigest(), CACHE_50_MINUTES)
|
|
cache.set(CLASSIFIER_VERSION_KEY, self.FORMAT_VERSION, CACHE_50_MINUTES)
|
|
return False
|
|
|
|
# subtract 1 since -1 (null) is also part of the classes.
|
|
|
|
# union with {-1} accounts for cases where all documents have
|
|
# correspondents and types assigned, so -1 isn't part of labels_x, which
|
|
# it usually is.
|
|
num_correspondents: int = len(set(labels_correspondent) | {-1}) - 1
|
|
num_document_types: int = len(set(labels_document_type) | {-1}) - 1
|
|
num_storage_paths: int = len(set(labels_storage_path) | {-1}) - 1
|
|
|
|
logger.debug(
|
|
f"{len(doc_pks)} documents, {num_tags} tag(s), {num_correspondents} correspondent(s), "
|
|
f"{num_document_types} document type(s). {num_storage_paths} storage path(s)",
|
|
)
|
|
|
|
from sklearn.feature_extraction.text import CountVectorizer
|
|
from sklearn.neural_network import MLPClassifier
|
|
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...")
|
|
|
|
def content_generator() -> Iterator[str]:
|
|
"""
|
|
Generates the content for documents, in the same order as the labels,
|
|
fetching it a chunk at a time
|
|
"""
|
|
for start in range(0, len(doc_pks), _CONTENT_CHUNK_SIZE):
|
|
chunk = doc_pks[start : start + _CONTENT_CHUNK_SIZE]
|
|
docs = Document.objects.only("content").order_by().in_bulk(chunk)
|
|
for pk in chunk:
|
|
# A document deleted since its labels were gathered still
|
|
# needs a row, so labels and content stay aligned
|
|
doc = docs.get(pk)
|
|
yield self.preprocess_content(
|
|
doc.content if doc is not None else "",
|
|
)
|
|
|
|
self.data_vectorizer = CountVectorizer(
|
|
analyzer="word",
|
|
ngram_range=(1, 2),
|
|
min_df=0.01,
|
|
)
|
|
|
|
data_vectorized: ndarray = self.data_vectorizer.fit_transform(
|
|
content_generator(),
|
|
)
|
|
|
|
# See the notes here:
|
|
# https://scikit-learn.org/stable/modules/generated/sklearn.feature_extraction.text.CountVectorizer.html
|
|
# This attribute isn't needed to function and can be large
|
|
self.data_vectorizer.stop_words_ = None
|
|
|
|
# Step 3: train the classifiers
|
|
if num_tags > 0:
|
|
logger.debug("Training tags classifier...")
|
|
notify(f"Training tags classifier ({num_tags} tag(s))...")
|
|
|
|
if num_tags == 1:
|
|
# Special case where only one tag has auto:
|
|
# Fallback to binary classification.
|
|
labels_tags = [
|
|
label[0] if len(label) == 1 else -1 for label in labels_tags
|
|
]
|
|
self.tags_binarizer = LabelBinarizer()
|
|
labels_tags_vectorized: ndarray = self.tags_binarizer.fit_transform(
|
|
labels_tags,
|
|
).ravel()
|
|
else:
|
|
self.tags_binarizer = MultiLabelBinarizer()
|
|
labels_tags_vectorized = self.tags_binarizer.fit_transform(labels_tags)
|
|
|
|
self.tags_classifier = MLPClassifier(tol=0.01, random_state=0)
|
|
self.tags_classifier.fit(data_vectorized, labels_tags_vectorized)
|
|
self._strip_training_state(self.tags_classifier)
|
|
else:
|
|
self.tags_classifier = None
|
|
logger.debug("There are no tags. Not training tags classifier.")
|
|
|
|
if num_correspondents > 0:
|
|
logger.debug("Training correspondent classifier...")
|
|
notify(
|
|
f"Training correspondent classifier ({num_correspondents} correspondent(s))...",
|
|
)
|
|
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),
|
|
)
|
|
self._strip_training_state(self.correspondent_classifier)
|
|
else:
|
|
self.correspondent_classifier = None
|
|
logger.debug(
|
|
"There are no correspondents. Not training correspondent classifier.",
|
|
)
|
|
|
|
if num_document_types > 0:
|
|
logger.debug("Training document type classifier...")
|
|
notify(
|
|
f"Training document type classifier ({num_document_types} type(s))...",
|
|
)
|
|
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),
|
|
)
|
|
self._strip_training_state(self.document_type_classifier)
|
|
else:
|
|
self.document_type_classifier = None
|
|
logger.debug(
|
|
"There are no document types. Not training document type classifier.",
|
|
)
|
|
|
|
if num_storage_paths > 0:
|
|
logger.debug(
|
|
"Training storage paths classifier...",
|
|
)
|
|
notify(f"Training storage path classifier ({num_storage_paths} path(s))...")
|
|
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),
|
|
)
|
|
self._strip_training_state(self.storage_path_classifier)
|
|
else:
|
|
self.storage_path_classifier = None
|
|
logger.debug(
|
|
"There are no storage paths. Not training storage path classifier.",
|
|
)
|
|
|
|
self.last_doc_change_time = latest_doc_change
|
|
self.last_auto_type_hash = hasher.digest()
|
|
self._update_data_vectorizer_hash()
|
|
|
|
# Set the classifier information into the cache
|
|
# Caching for 50 minutes, so slightly less than the normal retrain time
|
|
cache.set(CLASSIFIER_MODIFIED_KEY, self.last_doc_change_time, CACHE_50_MINUTES)
|
|
cache.set(CLASSIFIER_HASH_KEY, hasher.hexdigest(), CACHE_50_MINUTES)
|
|
cache.set(CLASSIFIER_VERSION_KEY, self.FORMAT_VERSION, CACHE_50_MINUTES)
|
|
|
|
return True
|
|
|
|
def preprocess_content(self, content: str) -> str:
|
|
"""
|
|
Process the contents of a document, distilling it down into
|
|
words which are meaningful to the content.
|
|
"""
|
|
language = settings.CLASSIFIER_LANGUAGE
|
|
content = unicodedata.normalize("NFC", content)
|
|
if language is None:
|
|
return " ".join(
|
|
match.group().lower() for match in RE_WORD.finditer(content)
|
|
)
|
|
return " ".join(_text_analyzer(language).analyze(content))
|
|
|
|
def _get_vectorizer_cache_key(self, content: str):
|
|
hash = sha256(content.encode())
|
|
hash.update(
|
|
f"|{self.FORMAT_VERSION}|{settings.CLASSIFIER_LANGUAGE}|{self.data_vectorizer_hash}".encode(),
|
|
)
|
|
return f"vectorized_content_{hash.hexdigest()}"
|
|
|
|
def _vectorize(self, content: str):
|
|
key = self._get_vectorizer_cache_key(content)
|
|
serialized_result = read_cache.get(key)
|
|
if serialized_result is None:
|
|
result = self.data_vectorizer.transform([self.preprocess_content(content)])
|
|
read_cache.set(key, signed_pickle_dumps(result), CACHE_5_MINUTES)
|
|
else:
|
|
try:
|
|
result = signed_pickle_loads(serialized_result)
|
|
except SignedPickleError:
|
|
result = self.data_vectorizer.transform(
|
|
[self.preprocess_content(content)],
|
|
)
|
|
read_cache.set(key, signed_pickle_dumps(result), CACHE_5_MINUTES)
|
|
else:
|
|
read_cache.touch(key, CACHE_5_MINUTES)
|
|
return result
|
|
|
|
def predict_correspondent(self, content: str) -> int | None:
|
|
if self.correspondent_classifier:
|
|
X = self._vectorize(content)
|
|
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)
|
|
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
|
|
|
|
if self.tags_classifier:
|
|
X = self._vectorize(content)
|
|
y = self.tags_classifier.predict(X)
|
|
tags_ids = self.tags_binarizer.inverse_transform(y)[0]
|
|
if type_of_target(y).startswith("multilabel"):
|
|
# the usual case when there are multiple tags.
|
|
return list(tags_ids)
|
|
elif type_of_target(y) == "binary" and tags_ids != -1:
|
|
# This is for when we have binary classification with only one
|
|
# tag and the result is to assign this tag.
|
|
return [tags_ids]
|
|
else:
|
|
# Usually binary as well with -1 as the result, but we're
|
|
# going to catch everything else here as well.
|
|
return []
|
|
else:
|
|
return []
|
|
|
|
def predict_storage_path(self, content: str) -> int | None:
|
|
if self.storage_path_classifier:
|
|
X = self._vectorize(content)
|
|
predicted_id = _predict_with_threshold(
|
|
self.storage_path_classifier,
|
|
X,
|
|
settings.CLASSIFIER_MATCH_THRESHOLD,
|
|
)
|
|
return predicted_id
|
|
return None
|