import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

model_id = "Qwen/Qwen2.5-Coder-0.5B-Instruct-GGUF"
model_file = "qwen2.5-coder-0.5b-instruct-q4_k_m.gguf"

print("Downloading model...")
model_path = hf_hub_download(repo_id=model_id, filename=model_file)
print("Loading model...")
llm = Llama(model_path=model_path, n_ctx=2048, n_threads=2)

def generate_code(prompt, max_tokens=512, temperature=0.7, top_p=0.9):
    messages = [
        {"role": "system", "content": "You are a helpful coding assistant."},
        {"role": "user", "content": prompt}
    ]
    output = llm.create_chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p
    )
    return output["choices"][0]["message"]["content"]

with gr.Blocks(title="Qwen 2.5 Coder 0.5B", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Qwen 2.5 Coder 0.5B")
    gr.Markdown("Powered by llama.cpp - Running locally on CPU")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Input",
                placeholder="Enter your coding question or prompt...",
                lines=5
            )
            with gr.Row():
                max_tokens = gr.Slider(minimum=64, maximum=2048, value=512, label="Max Tokens")
                temperature = gr.Slider(minimum=0.0, maximum=2.0, value=0.7, label="Temperature")
                top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.9, label="Top P")
            submit = gr.Button("Generate", variant="primary")
        with gr.Column():
            output = gr.Textbox(label="Output", lines=10)
    
    submit.click(
        fn=generate_code,
        inputs=[prompt, max_tokens, temperature, top_p],
        outputs=output
    )
    prompt.submit(
        fn=generate_code,
        inputs=[prompt, max_tokens, temperature, top_p],
        outputs=output
    )

demo.launch()