Spaces:

TeresaK
/

cpv_test

Runtime error

leavoigt commited on Jan 5

Commit

c6ebb8e

•

1 Parent(s): 69fa6c8

Update utils/target_classifier.py

Files changed (1) hide show

utils/target_classifier.py CHANGED Viewed

@@ -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)
     return doc_classifier

             classifier_name = config.get('target','MODEL')
     logging.info("Loading classifier")
+    # Loading classifier
+    doc_classifier = SetFitModel.from_pretrained("leavoigt/vulnerability_target")
     return doc_classifier