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)
    url = "https://drive.google.com/uc?id=1zi3hSBRTP9uwHDVwQMqpGeXb8RBhJfrO"
    model_path = gdown.download(url, quiet=False)
    model.load_weights(model_path)
    
    return model


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

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

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