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]))