from llama_cpp import Llama
import gradio as gr

llm = Llama(model_path="dcl2.gguf", n_ctx=256)

def generate(prompt, max_tokens):
    output = llm(prompt, max_tokens=int(max_tokens), stop=["</s>"])
    return output["choices"][0]["text"]

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Slider(1, 200, value=50, label="Max tokens"),
    ],
    outputs=gr.Textbox(label="Output"),
    title="dcl2 (custom GPT-2)",
)

demo.launch()