import streamlit as st
import tensorflow as tf
import numpy as np

# Load the trained model
model = tf.keras.models.load_model("news_classification_rnn.h5")

# Load Preprocessing Function
import dill
with open("preprocessing1.pkl", "rb") as f:
    clean_text = dill.load(f)

# Load Text Vectorization Layer from SavedModel
vectorizer = tf.saved_model.load("vectorizer")

# Define News Categories
news_categories = ["Business", "Sci/Tech", "Sports", "World"]

# Streamlit UI
st.title("📰 News Classification with Simple RNN")
st.write("Enter a news headline to predict its category.")

user_input = st.text_area("Enter News Text:", "")

if st.button("Classify"):
    if user_input.strip():
        # Preprocess input
        processed_text = clean_text(user_input)

        # Convert text to integer sequence
        text_sequence = vectorizer([processed_text])

        # Convert to numpy array (model expects batch input)
        text_sequence = np.array(text_sequence)

        # Predict Category
        prediction = model.predict(text_sequence)
        category = np.argmax(prediction)

        st.success(f"Predicted Category: **{news_categories[category]}**")
    else: 
        st.warning("⚠ Please enter a news headline.")