#TODO: verify seed works
#TODO: integrate into markdown
import gradio as gr
from gradio import mix
from transformers import pipeline, set_seed

#title = "trustworthy artificial intelligence workshop - content generator"
description = "based on the gpt2 demo interface by <a href='https://huggingface.co/spaces/docs-demos/gpt2/tree/main'>ahsen khaliq</a>"

#io1 = gr.Interface.load("huggingface/distilgpt2")

generator = pipeline('text-generation', model='gpt2')


io2 = gr.Interface.load("huggingface/gpt2-large")

#io3 = gr.Interface.load("huggingface/gpt2-medium")

#io4 = gr.Interface.load("huggingface/gpt2-xl")

def inference(text, seed):
    """
    if model == "gpt2-large":
        outtext = io2(text)
    elif model == "gpt2-medium":
        outtext = io3(text)
    elif model == "gpt2-xl":
        outtext = io4(text)
    else:
        outtext = io1(text)
    """
    #outtext = io2(text)
    set_seed(seed)
    outtext = generator(text, max_length=100, num_return_sequences=1)
    return outtext   
    
     

gr.Interface(
    inference, 
    [gr.inputs.Textbox(label="input", placeholder="trustworthy artificial intelligence"), gr.inputs.Number(label="seed")],
    #,gr.inputs.Dropdown(choices=["distilgpt2","gpt2-medium","gpt2-large","gpt2-xl"], type="value", default="gpt2-medium", label="model")], 
    gr.outputs.Textbox(label="gpt-2 proposal"),
    #title=title,
    #description=description,
    cache_examples=True).launch(enable_queue=True)

#TODO: add credits at bottom