Update utils/target_classifier.py

utils/target_classifier.py

@@ -74,8 +74,11 @@ def load_targetClassifier(config_file:str = None, classifier_name:str = None):
     
     logging.info("Loading classifier")  
 
+    # Get token from environment variable
+    hf_token = os.environ["model_token"]
+
     # Loading classifier
-    doc_classifier = SetFitModel.from_pretrained("GIZ/TARGET-VULNERABILITY-multiclass-mpnet")
+    doc_classifier = SetFitModel.from_pretrained(classifier_name, hf_token)
 
     return doc_classifier