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

model_path = hf_hub_download(
    repo_id="Kami574/Atlas-FirstAid",
    filename="Atlas-FirstAid.Q4_K_M.gguf"
)

llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=4,
    n_batch=512,
    n_gpu_layers=0,
    verbose=False,
)


@spaces.GPU(duration=60)
def chat_response(message, history):
    messages = []

    for msg in history:
        messages.append({
            "role": msg["role"],
            "content": msg["content"],
        })

    messages.append({
        "role": "user",
        "content": message,
    })

    response_stream = llm.create_chat_completion(
        messages=messages,
        temperature=0.7,
        max_tokens=384,
        stream=True
    )

    partial_message = ""
    for chunk in response_stream:
        if "choices" in chunk and chunk["choices"]:
            delta = chunk["choices"][0].get("delta", {})
            if "content" in delta:
                partial_message += delta["content"]
                yield partial_message


demo = gr.ChatInterface(
    fn=chat_response,
    title="Atlas - First Aid AI",
    description="Offline emergency guidance model.",
)


if __name__ == "__main__":
    demo.launch()