import os
from fastapi import FastAPI
import gradio as gr
from llama_cpp import Llama

# 1. Configuração do Modelo (Gemma 4 E4B GGUF)
# Usando uma versão quantizada para caber nos 16GB de RAM
model_id = "google/gemma-4-e4b-it-GGUF" 
model_file = "gemma-4-e4b-it-Q4_K_M.gguf"

# Inicializa o modelo (ele será baixado automaticamente se configurado no Space)
llm = Llama.from_pretrained(
    repo_id=model_id,
    filename=model_file,
    n_ctx=2048,           # Janela de contexto
    n_threads=2           # Limite de 2 vCPUs do Space gratuito
)

app = FastAPI()

def generate_response(message, history):
    # Formatação básica para o Gemma 4
    prompt = f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    
    output = llm(
        prompt,
        max_tokens=512,
        stop=["<|im_end|>"],
        echo=False
    )
    
    return output["choices"][0]["text"]

# 2. Interface Gradio
demo = gr.ChatInterface(
    fn=generate_response,
    title="Gemma 4 - E4B Thinking (CPU Free Tier)",
    description="Rodando Gemma 4 via GGUF no hardware gratuito da Hugging Face."
)

# 3. Montar Gradio dentro do FastAPI
app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)