import gradio as gr

iface = gr.Interface.load("models/nhantran0506/law-llms-v1")

def wrapper_fn(text, token):
    # Here, you can use the token as needed.
    # For example, you might pass it to your model.
    return iface.predict(text)

inputs = [
    gr.inputs.Textbox(lines=2, label="Input Text"),
    gr.inputs.Textbox(lines=1, label="Token")
]

iface = gr.Interface(fn=wrapper_fn, inputs=inputs, outputs="text")
iface.launch(share=True)