Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3
4
5device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
6tokenizer = AutoTokenizer.from_pretrained('Den4ikAI/ruBert-tiny-replicas-classifier')
7model = AutoModelForSequenceClassification.from_pretrained('Den4ikAI/ruBert-tiny-replicas-classifier')
8model.to(device)
9model.eval()
10
11classes = ['instruct', 'question', 'dialogue', 'problem', 'about_system', 'about_user']
12
13
14def get_sentence_type(text):
15 inputs = tokenizer(text, max_length=512, add_special_tokens=False, return_tensors='pt').to(device)
16 with torch.no_grad():
17 logits = model(**inputs).logits
18 probas = list(torch.sigmoid(logits)[0].cpu().detach().numpy())
19 out = classes[probas.index(max(probas))]
20 return out
21
22while 1:
23 print(get_sentence_type(input(":> ")))
24