Spaces:
Running
Running
Joshua Lochner
commited on
Commit
•
4678c9b
1
Parent(s):
74c1216
Add cache system for loading classifiers/vectorizers
Browse files- src/model.py +15 -6
src/model.py
CHANGED
@@ -101,11 +101,20 @@ def get_tokenizer(model_args, use_cache=True):
|
|
101 |
return tokenizer
|
102 |
|
103 |
|
104 |
-
|
105 |
-
|
106 |
-
|
107 |
-
|
108 |
-
|
109 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
110 |
|
111 |
return classifier, vectorizer
|
|
|
101 |
return tokenizer
|
102 |
|
103 |
|
104 |
+
CLASSIFIER_CACHE = {}
|
105 |
+
def get_classifier_vectorizer(classifier_args, use_cache=True):
|
106 |
+
classifier_path = os.path.join(classifier_args.classifier_dir, classifier_args.classifier_file)
|
107 |
+
if use_cache and classifier_path in CLASSIFIER_CACHE:
|
108 |
+
classifier = CLASSIFIER_CACHE[classifier_path]
|
109 |
+
else:
|
110 |
+
with open(classifier_path, 'rb') as fp:
|
111 |
+
classifier = CLASSIFIER_CACHE[classifier_path] = pickle.load(fp)
|
112 |
+
|
113 |
+
vectorizer_path = os.path.join(classifier_args.classifier_dir, classifier_args.vectorizer_file)
|
114 |
+
if use_cache and vectorizer_path in CLASSIFIER_CACHE:
|
115 |
+
vectorizer = CLASSIFIER_CACHE[vectorizer_path]
|
116 |
+
else:
|
117 |
+
with open(vectorizer_path, 'rb') as fp:
|
118 |
+
vectorizer = CLASSIFIER_CACHE[vectorizer_path] = pickle.load(fp)
|
119 |
|
120 |
return classifier, vectorizer
|