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

max_sequence_len = 40

st.title("Next Word Prediction using LSTM")
seed_text = st.text_input('Enter initial text to start generating next words', 'implementation of')
next_words_count = st.slider('How many next words to generate',  1, 8, 4)
model = load_model('text_generator.h5')

with open("tokenizer.pickle", "rb") as handle:
    tokenizer = pickle.load(handle)

def generate_text(seed_text, next_words, model, max_sequence_len):

    for _ in range(next_words):
        token_list = tokenizer.texts_to_sequences([seed_text])[0]

        token_list = pad_sequences([token_list], maxlen = max_sequence_len - 1, padding='pre')

        predicted = model.predict(token_list, verbose=0)

        classes_x = np.argmax(predicted,axis=1)

        output_word = ""
        for word,index in tokenizer.word_index.items():
            if index == classes_x:
                output_word = word
                break
        seed_text += " "+output_word
    return seed_text.title()

if st.button("Submit", type="primary"):
    output = generate_text(seed_text, next_words_count,model,max_sequence_len)
    st.write(output)