Joshua Lochner commited on
Commit
4678c9b
1 Parent(s): 74c1216

Add cache system for loading classifiers/vectorizers

Browse files
Files changed (1) hide show
  1. 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
- def get_classifier_vectorizer(classifier_args):
105
- with open(os.path.join(classifier_args.classifier_dir, classifier_args.classifier_file), 'rb') as fp:
106
- classifier = pickle.load(fp)
107
-
108
- with open(os.path.join(classifier_args.classifier_dir, classifier_args.vectorizer_file), 'rb') as fp:
109
- vectorizer = pickle.load(fp)
 
 
 
 
 
 
 
 
 
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