# install required packages
import subprocess
import sys

def install(package):
    subprocess.check_call([sys.executable, "-m", "pip", "install", package])

install("tensorflow")
install("numpy")
install("transformers")

# import related packages
import streamlit as st
import numpy as np
import tensorflow as tf

import transformers
from transformers import DistilBertTokenizer
from transformers import TFDistilBertForSequenceClassification

# print the header message
st.header("Welcome to the STEM NLP application!")

# fetch the pre-trained model
model = TFDistilBertForSequenceClassification.from_pretrained("kaixinwang/NLP")

# build the tokenizer
MODEL_NAME = 'distilbert-base-uncased'
# tokenizer = DistilBertTokenizer.from_pretrained(MODEL_NAME)
tokenizer = DistilBertTokenizer.from_pretrained("kaixinwang/NLP")

mapping = {0:"Negative", 1:"Positive"}
# prompt for the user input
x = st.text_input("To get started, enter your review/text below and hit ENTER:")
if x:
    st.write("Determining the sentiment...")
    # utterance tokenization
    encoding = tokenizer([x], truncation=True, padding=True)
    encoded = tf.data.Dataset.from_tensor_slices((dict(encoding), np.ones(1)))
    # make the prediction
    preds = model.predict(encoded.batch(1)).logits  
    prob = tf.nn.softmax(preds, axis=1).numpy()  
    prob_max = np.argmax(prob, axis=1)
    # display the output
    st.write("Your review is:", x)
    content = "Sentiment: %s, prediction score: %.4f" %(mapping[prob_max[0]], prob[0][prob_max][0])
    st.write(content)
    # st.write("Sentiment:", mapping[prob_max[0]], "Prediction Score:", prob[0][prob_max][0])