import streamlit as st
import joblib

st.title('SMS Spam Classification (Decision Tree)')

# Initialize model and vectorizer as None
model = None
vectorizer = None

# Load the model and vectorizer
try:
    model = joblib.load('src/decision_tree_model.joblib')
    vectorizer = joblib.load('src/vectorizer.joblib')  # Save your CountVectorizer as 'vectorizer.joblib'
except Exception as e:
    st.error(f"Error loading model or vectorizer: {e}")

text_input = st.text_area('Enter SMS text for classification:', '')

def preprocess_text(text, vectorizer):
    if hasattr(vectorizer, 'transform'):
        return vectorizer.transform([text])
    else:
        raise ValueError("Unknown vectorizer type.")

if st.button('Predict'):
    if not text_input.strip():
        st.warning('Please enter some text.')
    elif model is None or vectorizer is None:
        st.error('Model or vectorizer not loaded. Please check the files and try again.')
    else:
        try:
            X = preprocess_text(text_input, vectorizer)
            prediction = model.predict(X)[0]
            label = 'Spam' if prediction == 1 else 'Ham'
            st.success(f'Prediction: {label}')
        except Exception as e:
            st.error(f"Error making prediction: {e}")