import gradio as gr

# Load the model
model = gr.Interface.load("models/roneneldan/TinyStories-Instruct-33M")

# Create the interface
interface = gr.Interface(
    fn=model.predict,
    inputs=gr.Textbox(prompt="Type a prompt here..."),
    outputs=gr.Textbox(),
    live=True,
    capture_session=True
)

# Launch the interface
interface.launch()