13 lines
435 B
Python
13 lines
435 B
Python
from transformers import BertTokenizer, BertForSequenceClassification
|
|
|
|
# load tokenizer and model weights
|
|
tokenizer = BertTokenizer.from_pretrained('SkolkovoInstitute/russian_toxicity_classifier')
|
|
model = BertForSequenceClassification.from_pretrained('SkolkovoInstitute/russian_toxicity_classifier')
|
|
|
|
|
|
def detector(text):
|
|
# prepare the input
|
|
batch = tokenizer.encode(text, return_tensors='pt')
|
|
|
|
# inference
|
|
model(batch) |