Views
No views yet
1import torch
2import torch.nn.functional as F
3from transformers import BertForSequenceClassification, BertTokenizer
4from datasets import load_dataset
5
6model = BertForSequenceClassification.from_pretrained('boapps/kmdb_classification_model')
7tokenizer = BertTokenizer.from_pretrained('SZTAKI-HLT/hubert-base-cc')
8
9article = {'title': '400 milliós luxusvillába vette be magát Matolcsy és családja', 'description': 'Matolcsy György fiának cége megvette, Matolcsy György unokatestvérének bankja meghitelezte, Matolcsy György pedig használja a 430 millióért hirdetett II. kerületi luxusrezidenciát.'}
10
11tokenized_article = tokenizer(article['title']+'\n'+article['description'], return_tensors="pt")
12
13logits = model(**tokenized_article).logits
14probabilities = F.softmax(logits[0], dim=-1)
15
16print(probabilities)