Den4ikAI/ruBert_tiny_qa_detection

Эта модель предназначена для классификации диалоговых вопросов на QA и Dialog. Модель может применятся в различных чат-ботах для управления режимом диалога и интернет-поиска.

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
tokenizer = AutoTokenizer.from_pretrained('Den4ikAI/ruBert_tiny_qa_detection')
model = AutoModelForSequenceClassification.from_pretrained('Den4ikAI/ruBert_tiny_qa_detection')
model.to(device)
model.eval()


def classify(text):
    inputs = tokenizer(text, max_length=128, add_special_tokens=False, return_tensors='pt').to(device)
    with torch.no_grad():
        logits = model(**inputs).logits
        probas = torch.sigmoid(logits)[0].cpu().detach().numpy()
    return probas
print(classify(input(':> ')))