import gradio as gr
import json
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
vocab_size = 103
max_seq_len = 18
model = tf.keras.Sequential([
            tf.keras.layers.Embedding(vocab_size, 120, input_length=max_seq_len-1),
            tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(120)),
            tf.keras.layers.Dense(vocab_size, activation='softmax')
])
model.load_weights("model_weight/")
word_token = json.load("word_token.json")
# 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[predicted[0]]
        if "0" in output_word[-1]: 
            output_word = output_word[:-1]
            seed_text += " " + output_word +"."
            break
        seed_text += " " + output_word
    return seed_text
iface = gr.Interface(
    fn=text_generator, 
    inputs=["text","slider"], 
    outputs="text",
    title="Poem Bot",
    live=False
    )
iface.launch()