Views
No views yet
1import numpy as np
2import torch
3from transformers import AutoTokenizer, AutoModelForSequenceClassification
4
5# Prepare input texts. This model is pretrained on multi-lingual data
6# and fine-tuned on English
7test_texts = ['Utterance1;Utterance2;Utterance3']
8
9# Load the model and tokenizer
10model = AutoModelForSequenceClassification.from_pretrained(
11 'justtherightsize/small-e-czech-2stage-online-risks-cs', num_labels=5).to("cuda")
12
13tokenizer = AutoTokenizer.from_pretrained(
14 'justtherightsize/small-e-czech-2stage-online-risks-cs',
15 use_fast=False, truncation_side='left')
16assert tokenizer.truncation_side == 'left'
17
18# Define helper functions
19def predict_one(text: str, tok, mod, threshold=0.5):
20 encoding = tok(text, return_tensors="pt", truncation=True, padding=True,
21 max_length=256)
22 encoding = {k: v.to(mod.device) for k, v in encoding.items()}
23 outputs = mod(**encoding)
24 logits = outputs.logits
25 sigmoid = torch.nn.Sigmoid()
26 probs = sigmoid(logits.squeeze().cpu())
27 predictions = np.zeros(probs.shape)
28 predictions[np.where(probs >= threshold)] = 1
29 return predictions, probs
30
31def print_predictions(texts):
32 preds = [predict_one(tt, tokenizer, model) for tt in texts]
33 for c, p in preds:
34 print(f'{c}: {p.tolist():.4f}')
35
36# Run the prediction
37print_predictions(test_texts)