"""
Hugging Face Space for the Gemma 4 E2B QAT Q4_0 checkpoint.

Serves google/gemma-4-E2B-it-qat-q4_0-gguf with llama.cpp on CPU.
Q4 weights are much smaller than the bf16 fine-tune, and this path does not
wait for a ZeroGPU assignment.
"""

from __future__ import annotations

import os

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

QAT_REPO = os.environ.get("QAT_REPO") or "google/gemma-4-E2B-it-qat-q4_0-gguf"
GGUF_FILE = os.environ.get("GGUF_FILE") or "gemma-4-E2B_q4_0-it.gguf"

model_path = hf_hub_download(repo_id=QAT_REPO, filename=GGUF_FILE)
llm = Llama(
    model_path=model_path,
    n_ctx=1024,
    n_threads=2,
    n_gpu_layers=0,
    verbose=False,
)


def build_prompt(system: str, user: str) -> str:
    # Same framing as src/lib/model_client.js buildDriveThruPrompt().
    return "\n".join(
        [
            "<bos><start_of_turn>system",
            system,
            "<end_of_turn>",
            "<start_of_turn>user",
            user,
            "<end_of_turn>",
            "<start_of_turn>model",
            "",
        ]
    )


def generate(system: str, user: str, max_new_tokens: int = 120, temperature: float = 0.2) -> str:
    system = (system or "").strip()
    user = (user or "").strip()
    if not user:
        raise gr.Error("user text is required")

    max_new = max(1, min(int(max_new_tokens or 120), 256))
    temp = float(temperature if temperature is not None else 0.2)
    prompt = build_prompt(system, user)
    out = llm(
        prompt,
        max_tokens=max_new,
        temperature=temp,
        stop=["<end_of_turn>", "<eos>"],
    )
    text = out["choices"][0]["text"]
    return text.split("<end_of_turn>")[0].strip()


def health() -> str:
    return f"ok model={QAT_REPO} file={GGUF_FILE} quant=q4_0"


with gr.Blocks(title="Drive-thru Gemma QAT") as demo:
    gr.Markdown(
        "Gemma 4 E2B **QAT Q4_0** API. Weights stay in 4-bit, so this is faster "
        "than the bf16 checkpoint and does not use the ZeroGPU quota."
    )
    system = gr.Textbox(label="system", lines=2)
    user = gr.Textbox(label="user", lines=3)
    max_new_tokens = gr.Slider(8, 256, value=120, step=1, label="max_new_tokens")
    temperature = gr.Slider(0, 1, value=0.2, step=0.05, label="temperature")
    output = gr.Textbox(label="completion", lines=6)
    gr.Button("Generate").click(
        generate,
        [system, user, max_new_tokens, temperature],
        output,
        api_name="generate",
    )
    health_out = gr.Textbox(label="health")
    gr.Button("Health").click(health, None, health_out, api_name="health")


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