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):
    return txt + (res := llm(
        prompt=txt, 
        max_length=16384-len(txt),
        stop=['###', '///', '\n\n\n\n'],
    )['generated_text']), mdn+txt+res

code_input = gr.inputs.Code(label="Enter your code here")
submit_button = gr.outputs.Button(label="Submit")

options = [
    gr.outputs.Accordion([
        gr.outputs.Checkbox(label="Option 1"),
        gr.outputs.Checkbox(label="Option 2"),
        gr.outputs.Checkbox(label="Option 3")
    ], label="Options")
]
# COnvert this to gr.Blocks()

with gr.Blocks() as apps:
    with gr.Row():
        mdn = gr.Markdown("## StarCoder 15b Instruct\n\n---\n\n")
    with gr.Row():
        code = gr.Code("// Enter some Javascript here to interact with StarCoder.\n\nconsole.log(\"StarCoder says")
    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)