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

# Download GGUF directly from Hugging Face
model_path = hf_hub_download(
    repo_id="HauhauCS/Gemma4-12B-QAT-Uncensored-HauhauCS-Balanced",
    filename="Gemma4-12B-QAT-Uncensored-HauhauCS-Balanced-Q4_K_M.gguf"
)

# Load model using llama.cpp engine
llm = Llama(model_path=model_path, n_ctx=2048)

def generate(prompt, max_tokens=1024, temperature=0.6, top_p=0.9):
    output = llm(
        f"User: {prompt}\nAssistant:",
        max_tokens=int(max_tokens),
        temperature=temperature,
        top_p=top_p,
        stop=["User:"]
    )
    return output["choices"][0]["text"]

# Expose Gradio API interface
demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Slider(64, 2048, value=1024, label="Max Tokens"),
        gr.Slider(0.1, 1.0, value=0.6, label="Temperature"),
        gr.Slider(0.1, 1.0, value=0.9, label="Top P")
    ],
    outputs="text"
)

demo.launch()