Views
No views yet
1from transformers import AutoTokenizer, pipeline
2from transformers import ElectraConfig, ElectraForSequenceClassification
3
4tokenizer = AutoTokenizer.from_pretrained("deepset/electra-base-squad2")
5model = ElectraForSequenceClassification.from_pretrained('jquigl/electra-movie-genre')
6
7classifier = pipeline("text-classification", model = model, tokenizer = tokenizer)
8example = classifier("A Manny Velazquez underground slasher tribute about a number of serial killers on the loose in Chicago and begins a murder spree.")
9print(example[0])
10
11#LABEL_0 = Fantasy
12#LABEL_1 = Romance
13#LABEL_2 = Thriller
14#LABEL_3 = Biography
15#LABEL_4 = Horror
16#LABEL_5 = Action
17#LABEL_6 = Crime
18#LABEL_7 = Animation
19#LABEL_8 = Adventure
20#LABEL_9 = Mystery
21#LABEL_10 = War
22#LABEL_11 = Family
23#LABEL_12 = History
24#LABEL_13 = Scifi
25#LABEL_14 = Film-noir
26#LABEL_15 = Sports