Update utils/vulnerability_classifier.py
Browse files
utils/vulnerability_classifier.py
CHANGED
|
@@ -83,7 +83,7 @@ def load_vulnerabilityClassifier(config_file:str = None, classifier_name:str = N
|
|
| 83 |
hf_token = os.environ["model_token"]
|
| 84 |
|
| 85 |
# Download model from HF Hub
|
| 86 |
-
doc_classifier = SetFitModel.from_pretrained(classifier_name,
|
| 87 |
|
| 88 |
return doc_classifier
|
| 89 |
|
|
|
|
| 83 |
hf_token = os.environ["model_token"]
|
| 84 |
|
| 85 |
# Download model from HF Hub
|
| 86 |
+
doc_classifier = SetFitModel.from_pretrained(classifier_name, token=hf_token)
|
| 87 |
|
| 88 |
return doc_classifier
|
| 89 |
|