Views
No views yet
bert-base-uncased model was restricted to 1 hidden layer and
fine-tuned for sequence classification on the
imdb dataset loaded using the datasets library.1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2
3pretrained = "lannelin/bert-imdb-1hidden"
4
5tokenizer = AutoTokenizer.from_pretrained(pretrained)
6
7model = AutoModelForSequenceClassification.from_pretrained(pretrained)
8
9LABELS = ["negative", "positive"]
10
11def get_sentiment(text: str):
12 inputs = tokenizer.encode_plus(text, return_tensors='pt')
13
14 output = model(**inputs)[0].squeeze()
15
16 return LABELS[(output.argmax())]
17
18print(get_sentiment("What a terrible film!"))