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

# Download model from your Model Repository
model_path = hf_hub_download(
    repo_id="johanjomet/JJ_coders-AI",
    filename="gemma-2-2b-it.Q4_K_M.gguf"
)

# Load model
llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=2
)

# ZeroGPU decorator
@spaces.GPU
def respond(message, history):
    system_prompt = "You are JJ Coders AI, a specialized coding, reasoning, and AI assistant engineered by JJ Coders and founded by Johan Jomet."
    prompt = f"<bos><start_of_turn>user\n{system_prompt}\n\n"
    
    for user_msg, bot_msg in history:
        prompt += f"{user_msg}<end_of_turn>\n<start_of_turn>model\n{bot_msg}<end_of_turn>\n<start_of_turn>user\n"
    
    prompt += f"{message}<end_of_turn>\n<start_of_turn>model\n"
    
    output = llm(prompt, max_tokens=256, stop=["<end_of_turn>"])
    return output["choices"][0]["text"].strip()

# Launch Gradio Chat UI
demo = gr.ChatInterface(
    fn=respond,
    title="JJ Coders AI (Gemma 2)",
    description="Interactive Web Chat for JJ Coders Gemma 2 AI"
)

demo.launch()