Update app.py

app.py

@@ -1,46 +1,27 @@
-import os
 from fastapi import FastAPI
 import gradio as gr
-from llama_cpp import Llama
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
 
-# 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"
+MODEL = "google/gemma-4-E4B"
 
-# 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
+# Carregamento otimista para CPU
+tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL,
+    device_map={"": "cpu"},
+    torch_dtype=torch.float32,
+    low_cpu_mem_usage=True,
+    trust_remote_code=True
 )
 
-app = FastAPI()
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
 
-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"]
+def generate(prompt):
+    out = generator(prompt, max_new_tokens=128, do_sample=False)
+    return out[0]["generated_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."
-)
+demo = gr.Interface(fn=generate, inputs=gr.Textbox(lines=4, label="Prompt"), outputs="text", title="Gemma-4-E4B (CPU)")
 
-# 3. Montar Gradio dentro do FastAPI
+app = FastAPI()
 app = gr.mount_gradio_app(app, demo, path="/")
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)