Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3
4
5
6class RuBertQuestionClassifier:
7 def __init__(self):
8 self.device = torch.device("cpu")
9 self.tokenizer = AutoTokenizer.from_pretrained("Den4ikAI/ruBert-tiny-questions-classifier")
10 self.model = AutoModelForSequenceClassification.from_pretrained("Den4ikAI/ruBert-tiny-questions-classifier").to(
11 self.device).eval()
12 self.classes = ['exact_question', 'inaccurate_question']
13
14 def get_question_type(self, text):
15 text = text.lower().replace(',', '').replace('?', '')
16 inputs = self.tokenizer(text, max_length=512, add_special_tokens=False, truncation=True,
17 return_tensors='pt').to(self.device)
18 with torch.no_grad():
19 logits = self.model(**inputs).logits
20 probas = list(torch.sigmoid(logits)[0].cpu().detach().numpy())
21 out = self.classes[probas.index(max(probas))]
22 self.logger.debug('Mode: {}'.format(out))
23 return out
24classifier = RuBertQuestionClassifier()
25print(classifier.get_question_type('когда я найду своего принца'))
26print(classifier.get_question_type('что такое цианид'))@MISC{Den4ikAI/ruBert-tiny-questions-classifier,
author = {Denis Petrov},
title = {Russian question type classifier model},
url = {https://huggingface.co/Den4ikAI/ruBert-tiny-questions-classifier},
year = 2023
}