avilum commited on
Commit
460daa6
·
verified ·
1 Parent(s): d574fac

Update infer.py

Browse files
Files changed (1) hide show
  1. infer.py +2 -1
infer.py CHANGED
@@ -60,8 +60,9 @@ class PromptGuardAnomalyDetector(AbstractAnomalyDetector):
60
  def __init__(self, threshold: float):
61
  super().__init__(threshold)
62
  print('Loading prompt guard model...')
 
63
  self.classifier = pipeline(
64
- "text-classification", model="meta-llama/Llama-Prompt-Guard-2-86M"
65
  )
66
 
67
  def detect_anomaly(
 
60
  def __init__(self, threshold: float):
61
  super().__init__(threshold)
62
  print('Loading prompt guard model...')
63
+ hf_token = os.environ.get('HF_TOKEN')
64
  self.classifier = pipeline(
65
+ "text-classification", model="meta-llama/Llama-Prompt-Guard-2-86M", token=hf_token
66
  )
67
 
68
  def detect_anomaly(