import gradio as gr
import json
import pickle
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences

# Static constants
max_seq_len = 279
model = tf.keras.models.load_model('text_gen_weight.h5')

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


with open("word_token.json") as wt:
    word_token = json.load(wt)
# Main function
def text_generator(seed_text,n):

    for _ in range(n):
        token_list = tokenizer.texts_to_sequences([seed_text])[0]
        token_list = pad_sequences([token_list], maxlen=max_seq_len-1, padding='pre')
        predicted = np.argmax(model.predict(token_list), axis=-1)
        output_word = ""
        output_word = word_token.get(str(predicted[0]))
        seed_text += " " + output_word
    return seed_text
iface = gr.Interface(
    fn=text_generator, 
    inputs=["text","slider"], 
    outputs="text",
    title="Text Generator",
    live=False
    )
iface.launch(debug=True)