import streamlit as st
import torch
from typing import List
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
from transformers import TFAutoModelForSequenceClassification
from scipy.special import softmax

task = 'sentiment'
MODEL = f"cardiffnlp/twitter-roberta-base-sentiment"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForSequenceClassification.from_pretrained(MODEL)

def polarity_scores_roberta(example):
    encoded_text = tokenizer(example, truncation=True, return_tensors='pt')
    output = model(**encoded_text)
    scores = output[0][0].detach().numpy()
    scores = softmax(scores)
    scores_dict = {
        'Roberta_NEG' : scores[0],
        'Roberta_NEU' : scores[1],
        'Roberta_POS' : scores[2]
    }
    return scores_dict

text = st.text_area('enter some texts!')

if text:
    out = polarity_scores_roberta(text)
    st.json(out)