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
|