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

# ----------------------
# Load model & tokenizer
# ----------------------
@st.cache_resource
def load_assets():
    model = tf.keras.models.load_model("model.h5")  # your saved model
    with open("tokenizer.pkl", "rb") as f:                   # your saved tokenizer
        tokenizer = pickle.load(f)
    return model, tokenizer

model, tokenizer = load_assets()

# ----------------------
# Streamlit UI
# ----------------------
st.title("Team 8 Project Demo")
st.write("Type a sentence and let the model suggest the next word!")

# Input text
text = st.text_input("Enter your sentence:")

def predict_next_word(model, tokenizer, text, max_len=20):
    """Predict next word from input text using trained LSTM model."""
    seq = tokenizer.texts_to_sequences([text])[0]
    seq = tf.keras.preprocessing.sequence.pad_sequences([seq], maxlen=max_len-1, padding='pre')

    preds = model.predict(seq, verbose=0)[0]
    next_index = np.argmax(preds)
    for word, index in tokenizer.word_index.items():
        if index == next_index:
            return word
    return None

if st.button("Predict Next Word") and text:
    predicted_word = predict_next_word(model, tokenizer, text)
    if predicted_word:
        st.success(f"**Predicted next word:** {predicted_word}")
    else:
        st.warning("Could not predict a word. Try another input.")