Views
No views yet
1import numpy as np
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3
4# Prepare input texts. This model is fine-tuned for Czech
5test_texts = ['Utterance1;Utterance2;Utterance3']
6
7# Load the model and tokenizer
8model = AutoModelForSequenceClassification.from_pretrained(
9 'justtherightsize/small-e-czech-binary-online-risks-cs', num_labels=2).to("cuda")
10
11tokenizer = AutoTokenizer.from_pretrained(
12 'justtherightsize/small-e-czech-binary-online-risks-cs',
13 use_fast=False, truncation_side='left')
14assert tokenizer.truncation_side == 'left'
15
16# Define helper functions
17def get_probs(text, tokenizer, model):
18 inputs = tokenizer(text, padding=True, truncation=True, max_length=256,
19 return_tensors="pt").to("cuda")
20 outputs = model(**inputs)
21 return outputs[0].softmax(1)
22
23def preds2class(probs, threshold=0.5):
24 pclasses = np.zeros(probs.shape)
25 pclasses[np.where(probs >= threshold)] = 1
26 return pclasses.argmax(-1)
27
28def print_predictions(texts):
29 probabilities = [get_probs(
30 texts[i], tokenizer, model).cpu().detach().numpy()[0]
31 for i in range(len(texts))]
32 predicted_classes = preds2class(np.array(probabilities))
33 for c, p in zip(predicted_classes, probabilities):
34 print(f'{c}: {p}')
35
36# Run the prediction
37print_predictions(test_texts)