import gradio as gr
import loop

MODEL_ID = "TinyLlama/TinyLlama_v1.1_math_code"

handler = loop.EndpointHandler(MODEL_ID)

def respond(prompt, grammar):
    args = { "inputs": prompt, "grammar": grammar }
    return handler(prompt)

demo = gr.Interface(
    respond,
    inputs=["text", "text"],
    outputs=["text"]
)

if __name__ == "__main__":
    demo.launch()