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)
|