-
Notifications
You must be signed in to change notification settings - Fork 0
/
ejemplo.py
41 lines (31 loc) · 1.11 KB
/
ejemplo.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
EXAMPLE_TEXT = "texto_de_ejemplo"
MODEL_PATH = "hate_type_model"
TOKENIZER = 'pysentimiento/robertuito-base-uncased'
LABEL_MAPPING = {
0: 'general',
1: 'político',
2: 'misoginia',
3: 'xenofobia',
4: 'sexual'
}
def _model_prediction(model, tokenizer, text):
# Tokenize example text
tokens = tokenizer(text, return_tensors="pt")
# Predict
with torch.no_grad():
outputs = model(**tokens)
# Obtain each class probability
probabilidades = torch.softmax(outputs.logits, dim=1).squeeze()
# Obtain the higher class
clase_predicha = torch.argmax(probabilidades).item()
# Return prediction
return LABEL_MAPPING[clase_predicha]
# Load model and tokenizer
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER)
# Predict
prediction = _model_prediction(model, tokenizer, EXAMPLE_TEXT)
# Show prediction
print(f'El comentario:\n"{EXAMPLE_TEXT}"\n\nContiene odio del tipo: {prediction}')