import streamlit as st
import numpy as np
import pickle
from tensorflow.keras.models import load_model

model = load_model("model.h5")
with open("count_vec.pkl", "rb") as f:
    vectorizer = pickle.load(f)

def run():
    st.title("Prediction - Spam Message Detection Model")
    st.write("---")
    st.image('ham_or_spam.jpg')
    
    user_input = st.text_area("Enter the message to check for spam")

    if st.button("Predict"):
        if user_input:
            user_input_vectorized = vectorizer.transform([user_input])
            
            prediction = model.predict(user_input_vectorized)
            prediction_class = np.argmax(prediction, axis=1)[0]

            if prediction_class == 1:
                st.error("This message is predicted to be SPAM!")
            else:
                st.success("This message is NOT spam.")
        else:
            st.warning("Please enter a message to classify.")

if __name__ == '__main__':
    run()