Zekun Wu commited on
Commit
bfea8bb
·
1 Parent(s): a7e2fc8
stereotype_detector/stereotype_detector.py CHANGED
@@ -30,10 +30,10 @@ class Detector:
30
  # Using pipeline for inference
31
  if classifier == 'Token':
32
  task_type = "ner"
33
- self.model = pipeline(task_type, model=self.model_path, tokenizer=self.model_path,use_auth_token=API_token)
34
  else:
35
  task_type = "text-classification"
36
- self.model = pipeline(task_type, model=self.model_path, tokenizer=self.model_path,use_auth_token=API_token,return_all_scores= True)
37
 
38
 
39
  def predict(self, texts: List[str]):
 
30
  # Using pipeline for inference
31
  if classifier == 'Token':
32
  task_type = "ner"
33
+ self.model = pipeline(task_type, model=self.model_path, tokenizer=self.model_path,token=API_token)
34
  else:
35
  task_type = "text-classification"
36
+ self.model = pipeline(task_type, model=self.model_path, tokenizer=self.model_path,token=API_token,return_all_scores= True)
37
 
38
 
39
  def predict(self, texts: List[str]):