#app.py
from transformers import pipeline
import gradio as gr

# the text generation pipeline
pipe = pipeline("text-generation", model="kaiku03/wildchat2")

examples = [
    "her lips touches mine and she continues to",
    "In a world where magic is real,",
    "The sun sets over the horizon as",
]

# function for the gradio app
def fn_textgen(prompt):
    return pipe(prompt, max_length=100)[0]['generated_text']

# create a Gradio interface
gr.Interface(fn=fn_textgen, 
             inputs="textbox", 
             outputs=gr.Textbox(label="Generated text"),
             examples=examples,
             title="WildChat Text Generation", 
             description="Enter a prompt and watch WildChat generate a response!",
             article='This model was trained on 20% only of the wildchat dataset')

gr.launch()