File size: 811 Bytes
12788a7 d4115e6 12788a7 a1ee603 12788a7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | import sys
from transformers import BertTokenizer
from model import BertForMultiLabelClassification
from multilabel_pipeline import MultiLabelPipeline
#
# Run inference on text using geomotions model
# default_model stored with git-lfs -- make sure you have it installed
# as of now, default model == checkpt 5000
# This file can be imported _or_ run as a script with text as the first argument
#
m = "models/default_model"
tokenizer = BertTokenizer.from_pretrained(m)
model = BertForMultiLabelClassification.from_pretrained(m)
goemotions = MultiLabelPipeline(
model=model,
tokenizer=tokenizer,
threshold=0.1
)
def infer(text):
emo = goemotions([text])
emo = emo[0]
emo = dict(zip(emo['labels'], emo['scores']))
return emo
if __name__=="__main__":
print (infer(sys.argv[1]))
|