import gradio as gr
from transformers import pipeline
import keras_nlp
import tensorflow as tf
import gdown
from tensorflow import keras
import time

def download_model_NLP():
    preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
    "gpt2_base_en",
    sequence_length=128,)
    model = keras_nlp.models.GPT2CausalLM.from_preset(
    "gpt2_base_en", preprocessor=preprocessor)
    
    model_clear = model
    #output = "GPT2_keras_deer"
    #id_weights = "106fk9vGF_sHB8St9B2cKYE9jQVNSdKW_"
    #gdown.download(id=id_weights, output=output, quiet=False)

    #id_index = "10EWM2kJG4djpI5T6VLuj0UkGCbVIjmJx"
    #gdown.download(id=id_index, output=output, quiet=False)

    #id_check_point = "10Et7a0EzyrnnBhMWzt6frXLiNVm8a2qz"
    #gdown.download(id=id_check_point, output=output, quiet=False)

    #id_folder = "1zi3hSBRTP9uwHDVwQMqpGeXb8RBhJfrO"
    #gdown.download_folder(id=id_folder, quiet=True, use_cookies=False)

    # id = "1-CiBtOhzyVuiUGShONZqluey7TyX9-Y6"

    output = "total.h5"
    id = "1-KgcnP1ayWQ6l2-4h723JCYPoWxzOnU3"

    gdown.download(id=id, output=output, quiet=False)
    
    
    
    model.load_weights(output)
    #model = tf.keras.models.load_model(output)       
    
    return model_clear, model


def get_model():
    return pipeline('text-generation', model='gpt-2') 

def complete_text(start_of_sentence):

    result_clear = model_clear.generate(start_of_sentence, max_length=100)
    result = model.generate(start_of_sentence, max_length=100)
    #result = model(start_of_sentence, max_length=50, do_sample=True)[0]['generated_text']
    return result_clear, result

model_clear, model  = download_model_NLP()
iface = gr.Interface(fn=complete_text, inputs=gr.inputs.Textbox(lines=2, placeholder='Start of Sentence Here...'), outputs=["text", "text"],)
iface.launch()