import tensorflow as tf
import streamlit as st
from transformers import BertTokenizer
from Sastrawi.Stemmer.StemmerFactory import StemmerFactory
from transformers import TFBertForSequenceClassification


# Fungsi untuk memuat model BERT dan tokenizer
PRE_TRAINED_MODEL = 'indobenchmark/indobert-base-p2'
bert_tokenizer = BertTokenizer.from_pretrained(PRE_TRAINED_MODEL)
bert_model = TFBertForSequenceClassification.from_pretrained(PRE_TRAINED_MODEL, num_labels=2)
bert_model.load_weights('model.h5')

# Inisialisasi stemmer dari Sastrawi
stemmer = StemmerFactory().create_stemmer()  # Membuat stemmer Sastrawi

def preprocess_text(text):
    # Menggunakan Sastrawi untuk stemming
    stemmed_text = stemmer.stem(text.lower())

    return stemmed_text

def predict_sentiment(text):
    preprocessed_text = preprocess_text(text)  # Pra-pemrosesan teks dengan Sastrawi
    input_ids = tf.constant(bert_tokenizer.encode(preprocessed_text, add_special_tokens=True))[None, :]  # Menambahkan token khusus [CLS] dan [SEP]
    logits = bert_model(input_ids)[0]
    probabilities = tf.nn.softmax(logits, axis=1)
    sentiment = tf.argmax(probabilities, axis=1)
    return sentiment.numpy()[0]


# Judul aplikasi
st.title('Prediksi Sentimen menggunakan BERT')

# Input teks
text = st.text_area('Masukkan teks', '')

# Tombol untuk memprediksi sentimen
if st.button('Prediksi'):
    if text.strip() == '':
        st.warning('Masukkan teks terlebih dahulu.')
    else:
        sentiment, probabilities = predict_sentiment(text)
        
        # Menghitung persentase probabilitas sentimen positif
        positive_probability = probabilities[1] * 100
        negative_probability = probabilities[0] * 100
        st.write(f'HASIL PREDIKSI')
        if sentiment == 0:
            st.write(f'Negatif ({negative_probability:.2f}%)')
        else:
            st.write(f'Positif ({positive_probability:.2f}%)')