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

# Download model (adjust quantization as needed)
model_path = hf_hub_download(
    repo_id="HauhauCS/Qwen3.5-2B-Uncensored-HauhauCS-Aggressive",
    filename="Qwen3.5-2B-Uncensored-HauhauCS-Aggressive-Q4_K_M.gguf"
)

# Load model (adjust n_gpu_layers for GPU acceleration)
llm = Llama(
    model_path=model_path,
    n_ctx=4096,
    n_gpu_layers=-1,  # Offload all layers to GPU if available
    verbose=False
)

def chat(message, history):
    response = llm(
        f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n",
        max_tokens=512,
        stop=["<|im_end|>"],
        stream=True
    )
    output = ""
    for chunk in response:
        output += chunk["choices"][0]["text"]
        yield output

gr.ChatInterface(fn=chat, title="Qwen3.5-2B Uncensored").launch()