import streamlit as st
from streamlit.logger import get_logger
import pandas as pd 
from transformers import pipeline
from setfit import AbsaModel
import translators as ts
from config import base_model, senti_map, absa_model, senti_color
from annotated_text import annotated_text


LOGGER = get_logger(__name__)
#ts.preaccelerate_and_speedtest()
senti_task = pipeline("sentiment-analysis", model=base_model, tokenizer=base_model)
absa = AbsaModel.from_pretrained(absa_model[0], absa_model[1])

def run():
    st.write('Copy and paste comment into below text box.')
    txt = st.text_area('customer review')
    if len(txt.strip()) > 0:
        lang = st.selectbox('pick output language',  ['en', 'hi', 'zh'], index=0)
        with st.spinner(f'translate to {lang}'):
            txt_en = ts.translate_text(txt, to_language=lang, translator='google')
        
        with st.spinner('working on comment sentiment, please wait...'):
            sentiment = senti_task(txt_en)
        st.write(f"it's {senti_map[sentiment[0]['label']]} feedback with a confidence of {sentiment[0]['score']:.1%}")
        with st.spinner('detecting aspect sentiment...'):
            preds = absa(txt_en)
        st.write(f"the comment talks about: {','.join([t['span'] for t in preds])}, detailed sentiments as follow:")
        #st.write(f'Customer says: {txt_en}')
        preds = {p['span']: p['polarity'] for p in preds}
        annotated_text(
            [(t + ' ', preds[t], senti_color[preds[t]]) if t in preds else t+' ' for t in txt_en.split(' ')]
        )

if __name__ == "__main__":
    run()