Update utils/vulnerability_classifier.py

utils/vulnerability_classifier.py

@@ -8,6 +8,7 @@ from utils.preprocessing import processingpipeline
 import streamlit as st
 from transformers import pipeline
 from setfit import SetFitModel
+import os
 
 
 # labels dictionary
@@ -76,25 +77,14 @@ def load_vulnerabilityClassifier(config_file:str = None, classifier_name:str = N
             classifier_name = config.get('vulnerability','MODEL')
     
     logging.info("Loading vulnerability classifier")
-    
-    # we are using the pipeline as the model is multilabel and DocumentClassifier 
-    # from Haystack doesnt support multilabel
-    # in pipeline we use 'sigmoid' to explicitly tell pipeline to make it multilabel
-    # if not then it will automatically use softmax, which is not a desired thing.
-    # doc_classifier = TransformersDocumentClassifier(
-    #                     model_name_or_path=classifier_name,
-    #                     task="text-classification",
-    #                     top_k = None)
 
-    # Download model from HF Hub
-    doc_classifier = SetFitModel.from_pretrained(classifier_name)
+
+    # Get token from environment variable
+    hf_token = os.environ["model_token"]
     
+    # Download model from HF Hub
+    doc_classifier = SetFitModel.from_pretrained(classifier_name, model_token=hf_token)
     
-    # doc_classifier = pipeline("text-classification", 
-    #                         model=classifier_name, 
-    #                         return_all_scores=True, 
-    #                         function_to_apply= "sigmoid")
-
     return doc_classifier