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)
    
    
    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)

    print(output)
    model.load_weights(output)
    
    return model


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

def complete_text(start_of_sentence):
    
    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

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