Views
No views yet
1from transformers import AutoModelForSequenceClassification, TFAutoModelForSequenceClassification
2from transformers import AutoTokenizer
3import numpy as np
4from scipy.special import softmax
5
6
7MODEL = f"cardiffnlp/tweet-topic-19-single"
8tokenizer = AutoTokenizer.from_pretrained(MODEL)
9
10# PT
11model = AutoModelForSequenceClassification.from_pretrained(MODEL)
12class_mapping = model.config.id2label
13
14text = "Tesla stock is on the rise!"
15encoded_input = tokenizer(text, return_tensors='pt')
16output = model(**encoded_input)
17
18scores = output[0][0].detach().numpy()
19scores = softmax(scores)
20
21# TF
22#model = TFAutoModelForSequenceClassification.from_pretrained(MODEL)
23#class_mapping = model.config.id2label
24#text = "Tesla stock is on the rise!"
25#encoded_input = tokenizer(text, return_tensors='tf')
26#output = model(**encoded_input)
27#scores = output[0][0]
28#scores = softmax(scores)
29
30
31ranking = np.argsort(scores)
32ranking = ranking[::-1]
33for i in range(scores.shape[0]):
34 l = class_mapping[ranking[i]]
35 s = scores[ranking[i]]
36 print(f"{i+1}) {l} {np.round(float(s), 4)}")
371) business_&_entrepreneurs 0.8575
2) science_&_technology 0.0604
3) pop_culture 0.0295
4) daily_life 0.0217
5) sports_&_gaming 0.0154
6) arts_&_culture 0.0154