| ```python | |
| from transformers import BertTokenizer, BertForSequenceClassification | |
| import torch | |
| import numpy as np | |
| import json | |
| class Prehibition: | |
| def __init__(self): | |
| model_name = 'wyluilipe/prehibiton-themes-clf' | |
| self.tokenizer = BertTokenizer.from_pretrained(model_name) | |
| self.model = BertForSequenceClassification.from_pretrained(model_name) | |
| def predict(self, text): | |
| tokenized = self.tokenizer.batch_encode_plus( | |
| [text], | |
| max_length = 512, | |
| pad_to_max_length=True, | |
| truncation=True, | |
| return_token_type_ids=False | |
| ) | |
| tokens_ids, mask = torch.tensor(tokenized['input_ids']), torch.tensor(tokenized['attention_mask']) | |
| with torch.no_grad(): | |
| model_output = self.model(tokens_ids, mask) | |
| return np.argmax(model_output['logits']).item() | |
| ``` |