import gradio as gr
from llama_cpp import Llama

# Load the quantized model we built
# 'n_threads' is key for making CPUs move fast
llm = Llama(model_path="./my_custom_model_q4.gguf", n_ctx=2048, n_threads=2)

def generate_text(prompt):
    output = llm(f"User: {prompt}\nAI:", max_tokens=128, stop=["User:"], stream=True)
    response = ""
    for token in output:
        response += token['choices'][0]['text']
        yield response

demo = gr.Interface(fn=generate_text, inputs="text", outputs="text")
demo.launch()