app.py

"""
Kepler AI — Chat Interface (GGUF, 7B, Q5_K_M) on HF Spaces with ZeroGPU
Model: aryyanthakrr/Kepler-7B-Q5_K_M-GGUF
"""

import os

import gradio as gr
import spaces
from huggingface_hub import hf_hub_download, list_repo_files

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
MODEL_REPO = "aryyanthakrr/Kepler-70B-Orbit"
APP_TITLE = "Orbit AI"
APP_TAGLINE = "Your orbit around intelligence 🚀 (70B, Q4_K_M)"

SYSTEM_PROMPT = "You are Kepler AI, a helpful, precise, and friendly assistant."

QUICK_ACTIONS = {
    "💻 Write Code": "Write clean, well-commented code for: ",
    "🖼️ Image Prompt": "Write a detailed, vivid image-generation prompt for: ",
    "📄 Summarize": "Summarize the following text clearly and concisely:\n\n",
    "🧠 Explain": "Explain this in simple terms, as if to a beginner:\n\n",
}

HF_TOKEN = os.environ.get("HF_TOKEN")  # set in Space secrets if the repo is private

# ---------------------------------------------------------------------------
# Download the GGUF file once at startup (CPU-only work, no GPU needed here).
# ---------------------------------------------------------------------------
print(f"Looking up files in {MODEL_REPO} ...")
all_files = list_repo_files(MODEL_REPO, token=HF_TOKEN)
gguf_files = [f for f in all_files if f.endswith(".gguf")]

if not gguf_files:
    raise FileNotFoundError(f"No .gguf file found in {MODEL_REPO}. Files found: {all_files}")

gguf_filename = next((f for f in gguf_files if "q5_k_m" in f.lower()), gguf_files[0])
print(f"Using GGUF file: {gguf_filename}")

MODEL_PATH = hf_hub_download(repo_id=MODEL_REPO, filename=gguf_filename, token=HF_TOKEN)
print(f"Downloaded to: {MODEL_PATH}")


def _ensure_cuda_runtime_on_path():
    """
    llama-cpp-python's compiled .so needs libcudart.so.12 at import time.
    That library ships inside the `nvidia-cuda-runtime-cu12` pip package
    rather than being on the system path, so we point LD_LIBRARY_PATH at it
    before importing llama_cpp.
    """
    import site

    search_roots = list(site.getsitepackages())
    try:
        search_roots.append(site.getusersitepackages())
    except Exception:
        pass

    extra_paths = []
    for root in search_roots:
        candidate = os.path.join(root, "nvidia", "cuda_runtime", "lib")
        if os.path.isdir(candidate):
            extra_paths.append(candidate)
        candidate_cublas = os.path.join(root, "nvidia", "cublas", "lib")
        if os.path.isdir(candidate_cublas):
            extra_paths.append(candidate_cublas)

    if extra_paths:
        current = os.environ.get("LD_LIBRARY_PATH", "")
        os.environ["LD_LIBRARY_PATH"] = os.pathsep.join(extra_paths + [current])
        print(f"LD_LIBRARY_PATH updated with: {extra_paths}")
    else:
        print("Warning: could not find nvidia-cuda-runtime-cu12 libs on this system.")


# ---------------------------------------------------------------------------
# Lazy model loader — llama_cpp is imported AND the model is loaded only
# the first time a GPU is actually available, i.e. inside a @spaces.GPU call.
# Importing it at module level fails because the CUDA runtime isn't visible
# outside of a @spaces.GPU-decorated function.
# ---------------------------------------------------------------------------
_llm = None


def get_model():
    global _llm
    if _llm is None:
        _ensure_cuda_runtime_on_path()
        from llama_cpp import Llama

        print("Loading GGUF model onto GPU ...")
        _llm = Llama(
            model_path=MODEL_PATH,
            n_ctx=4096,
            n_gpu_layers=-1,  # offload all layers to GPU
            n_threads=8,
            verbose=False,
        )
        print("Model loaded.")
    return _llm


# ---------------------------------------------------------------------------
# Prompt building
# ---------------------------------------------------------------------------
def build_messages(history, user_message, attached_file_text):
    if attached_file_text:
        user_message = f"{user_message}\n\n[Attached file content]:\n{attached_file_text}"

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for turn in history:
        messages.append({"role": "user", "content": turn[0]})
        if turn[1]:
            messages.append({"role": "assistant", "content": turn[1]})
    messages.append({"role": "user", "content": user_message})
    return messages, user_message


# ---------------------------------------------------------------------------
# GPU-bound generation
# ---------------------------------------------------------------------------
@spaces.GPU(duration=60)
def generate_response(messages, max_new_tokens, temperature):
    llm = get_model()
    stream = llm.create_chat_completion(
        messages=messages,
        max_tokens=int(max_new_tokens),
        temperature=temperature,
        top_p=0.9,
        stream=True,
    )
    partial = ""
    for chunk in stream:
        delta = chunk["choices"][0]["delta"].get("content", "")
        if delta:
            partial += delta
            yield partial


def respond(user_message, history, attached_file_text, max_new_tokens, temperature):
    messages, user_message = build_messages(history, user_message, attached_file_text)
    history = history + [[user_message, ""]]

    for partial in generate_response(messages, max_new_tokens, temperature):
        history[-1][1] = partial
        yield history, ""


def apply_quick_action(action_label, current_text):
    template = QUICK_ACTIONS.get(action_label, "")
    return template + current_text


def read_attached_file(file_obj):
    if file_obj is None:
        return ""
    try:
        with open(file_obj.name, "r", encoding="utf-8", errors="ignore") as f:
            return f.read()[:4000]
    except Exception as e:
        return f"[Could not read file: {e}]"


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="blue", neutral_hue="slate")

custom_css = """
#kepler-header { text-align: center; padding: 12px 0 4px 0; }
#kepler-header h1 { font-size: 2rem; margin-bottom: 0; }
#kepler-header p { color: var(--body-text-color-subdued); margin-top: 2px; }
.quick-btn { min-width: 0 !important; }
"""

with gr.Blocks(title=APP_TITLE) as demo:
    gr.HTML(
        f"""
        <div id="kepler-header">
            <h1>🪐 {APP_TITLE}</h1>
            <p>{APP_TAGLINE} — powered by <code>{MODEL_REPO}</code></p>
        </div>
        """
    )

    attached_text_state = gr.State("")

    with gr.Row():
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(height=480, label="Kepler AI")

            with gr.Row():
                msg = gr.Textbox(
                    placeholder="Ask Kepler AI anything...",
                    show_label=False,
                    scale=5,
                    lines=2,
                )
                send_btn = gr.Button("Send", variant="primary", scale=1)

            with gr.Row():
                for label in QUICK_ACTIONS:
                    gr.Button(label, elem_classes="quick-btn").click(
                        apply_quick_action,
                        inputs=[gr.State(label), msg],
                        outputs=msg,
                    )

            with gr.Accordion("📎 Attach a file (optional)", open=False):
                file_upload = gr.File(label="Attach text/code file", file_types=[".txt", ".md", ".py", ".csv", ".json"])
                file_upload.change(read_attached_file, inputs=file_upload, outputs=attached_text_state)

            clear_btn = gr.Button("🗑️ Clear chat")

        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Settings")
            max_new_tokens = gr.Slider(64, 1024, value=384, step=64, label="Max new tokens")
            temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="Temperature")
            gr.Markdown("### 🎙️ Voice input\n*Coming soon — text chat only for now.*")
            gr.Markdown(
                f"""
                ---
                **Model:** `{MODEL_REPO}`
                **Format:** GGUF (Q5_K_M)
                **Engine:** llama.cpp
                **Hardware:** ZeroGPU
                """
            )

    send_btn.click(
        respond,
        inputs=[msg, chatbot, attached_text_state, max_new_tokens, temperature],
        outputs=[chatbot, msg],
    )
    msg.submit(
        respond,
        inputs=[msg, chatbot, attached_text_state, max_new_tokens, temperature],
        outputs=[chatbot, msg],
    )
    clear_btn.click(lambda: ([], "", ""), outputs=[chatbot, msg, attached_text_state])

if __name__ == "__main__":
    demo.queue().launch(theme=theme, css=custom_css)