Implemented the classifier model, including automatic tagging of new …

…documents

.gitignore

@@ -83,3 +83,6 @@ scripts/nuke
 
 # Static files collected by the collectstatic command
 static/
+
+# Classification Models
+models/

src/documents/admin.py

@@ -106,14 +106,6 @@ class CorrespondentAdmin(CommonAdmin):
     list_filter = ("matching_algorithm",)
     list_editable = ("match", "matching_algorithm")
 
-    def save_model(self, request, obj, form, change):
-        super().save_model(request, obj, form, change)
-
-        for document in Document.objects.filter(correspondent__isnull=True).exclude(tags__is_archived_tag=True):
-            if obj.matches(document.content):
-                document.correspondent = obj
-                document.save(update_fields=("correspondent",))
-
     def get_queryset(self, request):
         qs = super(CorrespondentAdmin, self).get_queryset(request)
         qs = qs.annotate(document_count=models.Count("documents"), last_correspondence=models.Max("documents__created"))
@@ -135,13 +127,6 @@ class TagAdmin(CommonAdmin):
     list_filter = ("colour", "matching_algorithm")
     list_editable = ("colour", "match", "matching_algorithm")
 
-    def save_model(self, request, obj, form, change):
-        super().save_model(request, obj, form, change)
-
-        for document in Document.objects.all().exclude(tags__is_archived_tag=True):
-            if obj.matches(document.content):
-                document.tags.add(obj)
-
     def get_queryset(self, request):
         qs = super(TagAdmin, self).get_queryset(request)
         qs = qs.annotate(document_count=models.Count("documents"))
@@ -158,14 +143,6 @@ class DocumentTypeAdmin(CommonAdmin):
     list_filter = ("matching_algorithm",)
     list_editable = ("match", "matching_algorithm")
 
-    def save_model(self, request, obj, form, change):
-        super().save_model(request, obj, form, change)
-
-        for document in Document.objects.filter(document_type__isnull=True).exclude(tags__is_archived_tag=True):
-            if obj.matches(document.content):
-                document.document_type = obj
-                document.save(update_fields=("document_type",))
-
     def get_queryset(self, request):
         qs = super(DocumentTypeAdmin, self).get_queryset(request)
         qs = qs.annotate(document_count=models.Count("documents"))

src/documents/apps.py

@@ -11,9 +11,7 @@ def ready(self):
         from .signals import document_consumption_started
         from .signals import document_consumption_finished
         from .signals.handlers import (
-            set_correspondent,
-            set_tags,
-            set_document_type,
+            classify_document,
             run_pre_consume_script,
             run_post_consume_script,
             cleanup_document_deletion,
@@ -22,9 +20,7 @@ def ready(self):
 
         document_consumption_started.connect(run_pre_consume_script)
 
-        document_consumption_finished.connect(set_tags)
-        document_consumption_finished.connect(set_correspondent)
-        document_consumption_finished.connect(set_document_type)
+        document_consumption_finished.connect(classify_document)
         document_consumption_finished.connect(set_log_entry)
         document_consumption_finished.connect(run_post_consume_script)
 

src/documents/classifier.py

@@ -0,0 +1,67 @@
+import pickle
+
+from documents.models import Correspondent, DocumentType, Tag
+from paperless import settings
+
+
+def preprocess_content(content):
+    content = content.lower()
+    content = content.strip()
+    content = content.replace("\n", " ")
+    content = content.replace("\r", " ")
+    while content.find("  ") > -1:
+        content = content.replace("  ", " ")
+    return content
+
+
+class DocumentClassifier(object):
+
+    @staticmethod
+    def load_classifier():
+        clf = DocumentClassifier()
+        clf.reload()
+        return clf
+
+    def reload(self):
+        with open(settings.MODEL_FILE, "rb") as f:
+            self.data_vectorizer = pickle.load(f)
+            self.tags_binarizer = pickle.load(f)
+            self.correspondent_binarizer = pickle.load(f)
+            self.type_binarizer = pickle.load(f)
+
+            self.tags_classifier = pickle.load(f)
+            self.correspondent_classifier = pickle.load(f)
+            self.type_classifier = pickle.load(f)
+
+    def save_classifier(self):
+        with open(settings.MODEL_FILE, "wb") as f:
+            pickle.dump(self.data_vectorizer, f)
+
+            pickle.dump(self.tags_binarizer, f)
+            pickle.dump(self.correspondent_binarizer, f)
+            pickle.dump(self.type_binarizer, f)
+
+            pickle.dump(self.tags_classifier, f)
+            pickle.dump(self.correspondent_classifier, f)
+            pickle.dump(self.type_classifier, f)
+
+    def classify_document(self, document, classify_correspondent=False, classify_type=False, classify_tags=False):
+        X = self.data_vectorizer.transform([preprocess_content(document.content)])
+
+        if classify_correspondent:
+            y_correspondent = self.correspondent_classifier.predict(X)
+            correspondent = self.correspondent_binarizer.inverse_transform(y_correspondent)[0]
+            print("Detected correspondent:", correspondent)
+            document.correspondent = Correspondent.objects.filter(name=correspondent).first()
+
+        if classify_type:
+            y_type = self.type_classifier.predict(X)
+            type = self.type_binarizer.inverse_transform(y_type)[0]
+            print("Detected document type:", type)
+            document.type = DocumentType.objects.filter(name=type).first()
+
+        if classify_tags:
+            y_tags = self.tags_classifier.predict(X)
+            tags = self.tags_binarizer.inverse_transform(y_tags)[0]
+            print("Detected tags:", tags)
+            document.tags.add(*[Tag.objects.filter(name=t).first() for t in tags])