#Back to Lesson 2, time flies!
from text_generation import Client
import gradio as gr

hf_api_key = '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'

#FalcomLM-instruct endpoint on the text_generation library
client = Client("https://api-inference.huggingface.co/models/tiiuae/falcon-40b-instruct", headers={"Authorization": f"Bearer {hf_api_key}"}, timeout=120)

def generate(input, slider):
    output = client.generate(input, max_new_tokens=slider).generated_text
    return output

demo = gr.Interface(fn=generate, inputs=[gr.Textbox(label="Prompt"), gr.Slider(label="Max new tokens", value=20,  maximum=1024, minimum=1)], outputs=[gr.Textbox(label="Completion")])
demo.launch()