Update utils/target_classifier.py

utils/target_classifier.py

@@ -68,10 +68,9 @@ def load_targetClassifier(config_file:str = None, classifier_name:str = None):
             classifier_name = config.get('target','MODEL')
     
     logging.info("Loading classifier")  
-      
-    doc_classifier = pipeline("text-classification", 
-                            model=classifier_name, 
-                            top_k =1)
+
+    # Loading classifier
+    doc_classifier = SetFitModel.from_pretrained("leavoigt/vulnerability_target")
 
     return doc_classifier