leavoigt commited on
Commit
c6ebb8e
1 Parent(s): 69fa6c8

Update utils/target_classifier.py

Browse files
Files changed (1) hide show
  1. utils/target_classifier.py +3 -4
utils/target_classifier.py CHANGED
@@ -68,10 +68,9 @@ def load_targetClassifier(config_file:str = None, classifier_name:str = None):
68
  classifier_name = config.get('target','MODEL')
69
 
70
  logging.info("Loading classifier")
71
-
72
- doc_classifier = pipeline("text-classification",
73
- model=classifier_name,
74
- top_k =1)
75
 
76
  return doc_classifier
77
 
 
68
  classifier_name = config.get('target','MODEL')
69
 
70
  logging.info("Loading classifier")
71
+
72
+ # Loading classifier
73
+ doc_classifier = SetFitModel.from_pretrained("leavoigt/vulnerability_target")
 
74
 
75
  return doc_classifier
76