Spam_Classification / prediction.py
SuryaR4421's picture
Upload 3 files
4f371bf verified
Raw
History Blame Contribute Delete
990 Bytes
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()