Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoConfig, MobileBertForSequenceClassification
3# load model
4model_name = r'cambridgeltl/sst_mobilebert-uncased'
5tokenizer = AutoTokenizer.from_pretrained(model_name)
6config = AutoConfig.from_pretrained(model_name)
7model = MobileBertForSequenceClassification.from_pretrained(model_name, config=config)
8model.eval()
9'''
10 labels:
11 0 -- negative
12 1 -- neutral
13 2 -- positive
14'''
15
16# prepare exemplar sentences
17batch_sentences = [
18 "in his first stab at the form , jacquot takes a slightly anarchic approach that works only sporadically .",
19 "a valueless kiddie paean to pro basketball underwritten by the nba .",
20 "a very well-made , funny and entertaining picture .",
21]
22
23# prepare input
24inputs = tokenizer(batch_sentences, max_length=256, truncation=True, padding=True, return_tensors='pt')
25input_ids, attention_mask = inputs.input_ids, inputs.attention_mask
26
27# make predictions
28outputs = model(input_ids=input_ids, attention_mask=attention_mask)
29predictions = torch.argmax(outputs.logits, dim = -1)
30print (predictions)
31# tensor([1, 0, 2])1@misc{susstmobilebert,
2 author = {Su, Yixuan},
3 title = {A MobileBERT Fine-tuned on SST},
4 howpublished = {\url{https://huggingface.co/cambridgeltl/sst_mobilebert-uncased}},
5 year = 2022
6}