Update utils/vulnerability_classifier.py

utils/vulnerability_classifier.py

@@ -7,6 +7,7 @@ from utils.config import getconfig
 from utils.preprocessing import processingpipeline
 import streamlit as st
 from transformers import pipeline
+from setfit import SetFitModel
 
 
 @st.cache_resource
@@ -42,10 +43,13 @@ def load_vulnerabilityClassifier(config_file:str = None, classifier_name:str = N
     #                     task="text-classification",
     #                     top_k = None)
 
-    doc_classifier = pipeline("text-classification", 
-                            model=classifier_name, 
-                            return_all_scores=True, 
-                            function_to_apply= "sigmoid")
+    # # Download model from HF Hub
+    doc_classifier = SetFitModel.from_pretrained("leavoigt/vulnerable_groups")
+    
+    # doc_classifier = pipeline("text-classification", 
+    #                         model=classifier_name, 
+    #                         return_all_scores=True, 
+    #                         function_to_apply= "sigmoid")
 
     return doc_classifier