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

print("Downloading model weights directly into Space storage...")
# We keep using Llama-3.2-3B because it fits beautifully in the 16GB free RAM limit
model_path = hf_hub_download(
    repo_id="unsloth/Llama-3.2-3B-Instruct-GGUF", 
    filename="Llama-3.2-3B-Instruct-Q4_K_M.gguf"
)

print("Loading model into CPU RAM...")
# CRITICAL: n_gpu_layers=0 tells llama.cpp NOT to look for a graphics card
llm = Llama(model_path=model_path, n_ctx=2048, n_gpu_layers=0)

def local_chat_stream(message, history):
    formatted_prompt = ""
    for user_msg, bot_msg in history:
        formatted_prompt += f"<|start_header_id|>user<|end_header_id|>\n{user_msg}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n{bot_msg}<|eot_id|>"
    formatted_prompt += f"<|start_header_id|>user<|end_header_id|>\n{message}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n"
    
    output_stream = llm(formatted_prompt, max_tokens=256, temperature=0.7, stream=True)
    
    partial_text = ""
    for chunk in output_stream:
        partial_text += chunk['choices'][0]['text']
        yield partial_text

gr.ChatInterface(
    fn=local_chat_stream,
    title="My Free CPU Chatbot App"
).launch()