| import streamlit as st | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TextClassificationPipeline | |
| import operator | |
| def get_sentiment(out): | |
| d = dict() | |
| for k in out.keys(): | |
| label = out[k]['label'] | |
| score = out[k]['score'] | |
| d[label] = score | |
| winning_lab = max(d.iteritems(), key=operator.itemgetter(1))[0] | |
| winning_score = d[winning_lab] | |
| return winning_lab, winning_score | |
| neg = out[0] | |
| neu = out[1] | |
| pos = out[2] | |
| model_name = "mrm8488/distilroberta-finetuned-financial-news-sentiment-analysis" | |
| model = AutoModelForSequenceClassification.from_pretrained(model_name) | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| pipe = TextClassificationPipeline(model=model, tokenizer=tokenizer, return_all_scores=True) | |
| text = st.text_area(f'Ciao! This app uses {model_name}.\nEnter your text to test it ❤️') | |
| if text: | |
| out = pipe(text) | |
| st.json(get_sentiment(out[0]) | |