import gradio as gr
from llama_cpp import Llama
import os

# Initialize the model directly from your repository
# This method handles the download and caching automatically
llm = Llama.from_pretrained(
    repo_id="harishkumar12k/Qwen2.5-Coder-1.5B-Q4_K_M-GGUF",
    filename="qwen2.5-coder-1.5b-q4_k_m.gguf",
    n_ctx=2048,
    n_threads=2 # Free tier typically has 2 vCPUs
)

def chat(message, history):
    # Standard ChatML format often used by Qwen
    prompt = "system\nYou are a helpful coding assistant.\n"
    for user_msg, assistant_msg in history:
        prompt += f"user\n{user_msg}\nassistant\n{assistant_msg}\n"
    prompt += f"user\n{message}\nassistant\n"

    output = llm(
        prompt,
        max_tokens=512,
        stop=["user", "\nuser", "<|endoftext|>"],
        echo=False
    )
    
    return output["choices"][0]["text"].strip()

demo = gr.ChatInterface(
    fn=chat,
    title="Qwen 2.5 Coder 1.5B (GGUF)",
    description="Custom GGUF running on Hugging Face Spaces CPU."
)

if __name__ == "__main__":
    # Must use 0.0.0.0 and 7860 for Docker Spaces
    demo.launch(server_name="0.0.0.0", server_port=7860)