import gradio as gr
from llama_cpp import Llama

llm = Llama.from_pretrained("bartowski/starcoder2-15b-instruct-v0.1-GGUF", filename="starcoder2-15b-instruct-v0.1-Q6_K.gguf", n_ctx=16384, n_gpu_layers=99)

def code_block(txt, mdn):
    mdn = mdn + "\n\n```javascript\n"
    res = llm(
        prompt=txt,
        stop=['###', '///', '\n\n\n'],
        stream=True
    )
    for chk in res:
        txt += chk['choices'][0]['text']
        yield txt, mdn + txt + "\n````\n"
    return txt, mdn+txt+"\n```\n"
    

with gr.Blocks() as apps:
    with gr.Row():
        mdn = gr.Markdown("## StarCoder 15b Instruct\n\n---\n\n")
    with gr.Row():
        code = gr.Code("// code.js\n/// Write a function in Javascript that queries a websocket for the current time from epoch in ms.\nconst epochy = () => {")
    with gr.Row():
        btn = gr.Button("Complete with StarCoder")
    btn.click(code_block, inputs=[code, mdn], outputs=[code, mdn])
apps.launch(share=False, debug=True)