leavoigt commited on
Commit
3c6fce1
1 Parent(s): 185893f

Update utils/vulnerability_classifier.py

Browse files
Files changed (1) hide show
  1. utils/vulnerability_classifier.py +1 -1
utils/vulnerability_classifier.py CHANGED
@@ -88,7 +88,7 @@ def load_vulnerabilityClassifier(config_file:str = None, classifier_name:str = N
88
  # top_k = None)
89
 
90
  # # Download model from HF Hub
91
- doc_classifier = SetFitModel.from_pretrained("leavoigt/vulnerable_groups")
92
 
93
  # doc_classifier = pipeline("text-classification",
94
  # model=classifier_name,
 
88
  # top_k = None)
89
 
90
  # # Download model from HF Hub
91
+ doc_classifier = SetFitModel.from_pretrained("leavoigt/vulnerability_multilabel")
92
 
93
  # doc_classifier = pipeline("text-classification",
94
  # model=classifier_name,