import streamlit as st
from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
import tensorflow as tf

tokenizer = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
model = TFAutoModelForSequenceClassification.from_pretrained('spectre0108/roberta-finetune-slangs')

# Function to predict sentiment
def predict_sentiment(text):
    tokenized = tokenizer(text, padding=True, truncation=True, return_tensors='tf', max_length=50)
    output = model(tokenized)
    logits = output.logits.numpy()

    probabilities = tf.nn.softmax(logits, axis=1).numpy()

    predicted_label = tf.argmax(probabilities, axis=1).numpy().item()
    positive_prob = probabilities[0][1]
    negative_prob = probabilities[0][0]

    return positive_prob, negative_prob

# Streamlit app
st.title("Sentiment Analysis for slangs with Fine-Tuned Transformer Model")

# Input text box
text_input = st.text_input("Enter a sentence:")

if st.button("Predict"):
    if text_input:
        positive_prob, negative_prob = predict_sentiment(text_input)
        st.success(f"Positive Probability: {positive_prob:.4f}")
        st.success(f"Negative Probability: {negative_prob:.4f}")
    else:
        st.warning("Please enter a sentence.")