from tensorflow.keras.models import load_model
from pickle import load
import numpy as np
import tensorflow as tf
import gradio as gr


# Set the model to the saved trained 300 epoch model. 
model = load_model('four_chapters_moby_dick_model_300_FIRAS.keras')
# Set the tokenizer to the trained tokenizer from the model. 
tokenizer = load(open('four_chapters_moby_dick_tokenizer_300_FIRAS', 'rb'))

def preprocess(texts):
    X = np.array(tokenizer.texts_to_sequences([texts])) -1
    return X
  
def next_word(text, num_gen_words=0, 
              randome_sampling = False,
              temperature=1):  
    '''
    Author : Firas Obeid
    Randome_Sampling : Using a categorical distribution to predict the character returned by the model
    Low temperatures results in more predictable text.
    Higher temperatures results in more surprising text.
    Experiment to find the best setting.
    '''
    input_text = text  
    output_text = []  
      
    for i in range(num_gen_words):  
        X_new = preprocess(input_text)  
        
        if randome_sampling:
            y_proba = model.predict(X_new, verbose = 0)[0, -1:, :]#first sentence, last token
            rescaled_logits = tf.math.log(y_proba) / temperature
            pred_word_ind = tf.random.categorical(rescaled_logits, num_samples=1) + 1
            pred_word = tokenizer.sequences_to_texts(pred_word_ind.numpy())[0]
        else:
            y_proba = model.predict(X_new, verbose=0)[0]  #first sentence
            pred_word_ind = np.argmax(y_proba, axis = -1) +1 
            pred_word = tokenizer.index_word[pred_word_ind[-1]]   
             
        
        input_text += ' ' + pred_word  
        output_text.append(pred_word)  
      
    return ' '.join(output_text)  
  
def generate_text(text, num_gen_words=25, temperature=1, randome_sampling=False):  
    return next_word(text, num_gen_words, randome_sampling, temperature)  

# Create an instance of the Gradio Interface application function with the appropriate parameters. 
max_output = gr.Number(value=150)
app = gr.Interface(fn=generate_text, 
                   inputs=["text", 
                           gr.Slider(1, 50, value=1, step=1, label="Minimum number of Shakespearean words to generate", info="Choose between 1 and 50"),
                           gr.Slider(0.1, 5, value=0.1, step=0.1, label="Temprature", info="Choose between 0.1 and 5"),
                          "checkbox"],
                   outputs="text")
# Launch the app
if __name__ == '__main__':
    app.launch(share=True)