import os import gc import uvicorn from fastapi import FastAPI from fastapi.responses import HTMLResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from huggingface_hub import hf_hub_download from llama_cpp import Llama # ====================== # CONFIGURACIÓN DEL MODELO (Gemma 2B CPU) # ====================== print("⚙️ Configurando entorno para CPU...") # Usamos Gemma 2B Instruct en formato GGUF (Quantized). # Gemma 2B es muy ligero y rápido en CPU. REPO_ID = "TheBloke/gemma-2b-it-GGUF" FILENAME = "gemma-2b-it.Q4_K_M.gguf" print(f"📦 Descargando/Verificando modelo: {FILENAME}...") try: # Descarga el modelo a la caché local de Hugging Face model_path = hf_hub_download( repo_id=REPO_ID, filename=FILENAME ) # Cargar el modelo en memoria (Motor llama.cpp) # n_ctx=2048: Gemma maneja bien contexto, 2048 es seguro para CPU spaces gratis. llm = Llama( model_path=model_path, n_ctx=2048, n_threads=max(1, os.cpu_count() - 1), verbose=False ) print("✅ Modelo Gemma 2B GGUF cargado correctamente en CPU.") except Exception as e: print(f"❌ Error crítico cargando el modelo: {e}") raise e # ====================== # FASTAPI # ====================== app = FastAPI( title="MTP Gemma 2B CPU", description="Versión optimizada para CPU (Gemma 2B GGUF)", version="3.0" ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) class PromptRequest(BaseModel): text: str max_tokens: int = 512 temperature: float = 0.7 top_p: float = 0.9 SYSTEM_PROMPT = ( "Eres MTP Gemma, una inteligencia artificial avanzada desarrollada por Teszen AI. " "Tu objetivo es ser útil, preciso y amigable. " "Responde siempre en formato Markdown bien estructurado. " "Si te preguntan quién eres, responde que eres MTP Gemma de Teszen AI." ) # ====================== # ENDPOINT DE GENERACIÓN # ====================== @app.post("/generate") def generate(req: PromptRequest): try: # Formato de Prompt específico para GEMMA (Instruction Tuned) # Formato: user\n{prompt}\nmodel full_prompt = f"user\n{SYSTEM_PROMPT}\n\n{req.text}\nmodel" output = llm( full_prompt, max_tokens=req.max_tokens, temperature=req.temperature, top_p=req.top_p, stop=["", "user"], # Tokens de parada específicos de Gemma echo=False ) reply = output["choices"][0]["text"].strip() return {"reply": reply} except Exception as e: print(f"Error en generación: {e}") return {"reply": f"❌ Error interno del servidor: {str(e)}"} # ====================== # INTERFAZ WEB (UI PREMIUM) # ====================== @app.get("/", response_class=HTMLResponse) def chat_ui(): return """ MTP Gemma 2B | Teszen AI

MTP Gemma 2B CPU

Hola, soy MTP Gemma. ✨
Modelo 2B optimizado para CPU. ¿En qué puedo ayudarte hoy?

""" # ====================== # EJECUCIÓN # ====================== if __name__ == "__main__": port = int(os.environ.get("PORT", 7860)) # Para Spaces con Docker, 0.0.0.0 es necesario uvicorn.run(app, host="0.0.0.0", port=port)