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

MODEL_REPO = "unsloth/Qwen3-VL-2B-Instruct-GGUF"
MODEL_FILE = "Qwen3-VL-2B-Instruct-UD-Q8_K_XL.gguf"

HF_TOKEN = os.getenv("HF_TOKEN")

model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
    token=HF_TOKEN
)

llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=2,
    n_gpu_layers=0,
    n_batch=128,
    verbose=False
)

SYSTEM_PROMPT = """És um assistente útil, claro e directo. Responde em português quando o utilizador escrever em português."""

def respond(message, history):
    prompt = f"System: {SYSTEM_PROMPT}\n\n"

    for item in history:
        if item["role"] == "user":
            prompt += f"User: {item['content']}\n"
        elif item["role"] == "assistant":
            prompt += f"Assistant: {item['content']}\n"

    prompt += f"User: {message}\nAssistant:"

    output = llm(
        prompt,
        max_tokens=256,
        temperature=0.6,
        top_p=0.9,
        stop=["User:", "System:"]
    )

    return output["choices"][0]["text"].strip()

demo = gr.ChatInterface(
    fn=respond,
    type="messages",
    title="Gemma 4 E4B - CPU Basic",
    description="Demo em Hugging Face Spaces CPU Basic usando GGUF + llama.cpp."
)

demo.launch()