Tbruand commited on
Commit
57b86c9
·
1 Parent(s): fb6559c

feat(model): ajoute le modèle CamemBERT fine-tuné avec renommage des labels

Browse files
Files changed (1) hide show
  1. models/fine_tuned.py +22 -0
models/fine_tuned.py CHANGED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import CamembertTokenizer, CamembertForSequenceClassification
2
+ import torch
3
+
4
+ class FineTunedModel:
5
+ def __init__(self):
6
+ model_id = "ymokay/toxicheck-camembert"
7
+ self.tokenizer = CamembertTokenizer.from_pretrained(model_id)
8
+ self.model = CamembertForSequenceClassification.from_pretrained(model_id)
9
+ self.model.eval()
10
+ self.label_map = {
11
+ "LABEL_0": "non-toxique",
12
+ "LABEL_1": "toxique"
13
+ }
14
+
15
+ def predict(self, text: str):
16
+ inputs = self.tokenizer(text, return_tensors="pt", truncation=True, padding=True)
17
+ with torch.no_grad():
18
+ logits = self.model(**inputs).logits
19
+ probs = torch.softmax(logits, dim=1).squeeze()
20
+
21
+ labels = [self.label_map.get(self.model.config.id2label[i], self.model.config.id2label[i]) for i in range(len(probs))]
22
+ return [(label, float(probs[i])) for i, label in enumerate(labels)]