paradizelost/paperless-ngx
1
0
Fork 0
mirror of https://github.com/paperless-ngx/paperless-ngx.git synced 2025-10-30 03:56:23 -05:00
Code Issues Packages Projects Releases Wiki Activity

write classifier model to temporary file before copying to final location

Browse Source
This commit is contained in:
jonaswinkler
2021-06-13 12:03:20 +02:00
parent 1ec54501c3
commit ddd9ac9a07
1 changed files with 9 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/documents/classifier.py
View File

@@ -3,6 +3,7 @@ import logging
import os import os
import pickle import pickle
import re import re
import shutil
from django.conf import settings from django.conf import settings
@@ -96,7 +97,10 @@ class DocumentClassifier(object):
raise ClassifierModelCorruptError() raise ClassifierModelCorruptError()
def save(self): def save(self):
with open(settings.MODEL_FILE, "wb") as f: target_file = settings.MODEL_FILE
target_file_temp = settings.MODEL_FILE + ".part"
with open(target_file_temp, "wb") as f:
pickle.dump(self.FORMAT_VERSION, f) pickle.dump(self.FORMAT_VERSION, f)
pickle.dump(self.data_hash, f) pickle.dump(self.data_hash, f)
pickle.dump(self.data_vectorizer, f) pickle.dump(self.data_vectorizer, f)
@@ -107,6 +111,10 @@ class DocumentClassifier(object):
pickle.dump(self.correspondent_classifier, f) pickle.dump(self.correspondent_classifier, f)
pickle.dump(self.document_type_classifier, f) pickle.dump(self.document_type_classifier, f)
if os.path.isfile(target_file):
os.unlink(target_file)
shutil.move(target_file_temp, target_file)
def train(self): def train(self):
data = list() data = list()