import os
import gradio as gr
from llama_cpp import Llama

# Configure llama.cpp to stream directly from your GGUF repository card
print("Loading model from Hub...")
llm = Llama.from_pretrained(
    repo_id="micymike/codemate-qwen-1.5B-8k-GGUF",
    filename="*q4_k_m.gguf", # Pulls your standard 4-bit optimized GGUF file
    n_ctx=4096,              # Bounded to 4k for the free CPU tier to maximize speed
    n_threads=2              # Fits cleanly into the free HF Spaces 2 vCPU limitation
)
print("Model loaded successfully!")

def generate(prompt, max_new_tokens=1024, temperature=0.7, top_p=0.9):
    # Enforces your fine-tuned chat format natively
    formatted_prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
    
    response = llm(
        formatted_prompt,
        max_tokens=int(max_new_tokens),
        temperature=float(temperature),
        top_p=float(top_p),
        stop=["<|im_end|>"] # Hard stop token to prevent generation loops
    )
    
    return response["choices"]["text"]

# Set up the frontend UI interface layout
demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt", lines=6, placeholder="Paste code or ask a debugging question here..."),
        gr.Slider(64, 2048, value=512, step=64, label="Max new tokens"),
        gr.Slider(0.1, 1.2, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p"),
    ],
    outputs=gr.Textbox(label="CodeMate Response", lines=14),
    title="🤖 CodeMate Qwen Playground",
    description="An optimized multi-turn Python & ReactJS conversational debugging assistant by Michael Moses.",
)

# Crucial for Hugging Face Spaces: Bind container to route traffic globally on port 7860
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)