Cricles commited on
Commit
fa61f7e
·
verified ·
1 Parent(s): 9375b5c

Update FRIDA/model.py

Browse files
Files changed (1) hide show
  1. FRIDA/model.py +1 -1
FRIDA/model.py CHANGED
@@ -102,7 +102,7 @@ class ToxicityPredictionResponse(BaseModel):
102
  predictions: List[ToxicityPrediction]
103
 
104
 
105
- def generate_resp(texts: List[str]):
106
  probs = infer(model, tokenizer, texts, device)
107
  probs_arr = probs.to('cpu').numpy()
108
  predictions = torch.argmax(probs, dim=-1).int().to('cpu').numpy()
 
102
  predictions: List[ToxicityPrediction]
103
 
104
 
105
+ def generate_resp(texts: List[str],model, tokenizer):
106
  probs = infer(model, tokenizer, texts, device)
107
  probs_arr = probs.to('cpu').numpy()
108
  predictions = torch.argmax(probs, dim=-1).int().to('cpu').numpy()