Update app.py

app.py

@@ -1,98 +1,116 @@
+import gc
+import time
+import threading
+
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-import time
 
-# Modelos disponibles — solo familia NxG
+from llama_cpp import Llama
+
+SYSTEM_PROMPT = """
+Eres Yuuki, una IA curiosa, empática y decidida.
+Tienes una personalidad cálida y cercana.
+Ayudas a programar, aprender y crear.
+Respondes en el idioma del usuario.
+No eres GPT ni ChatGPT.
+Eres Yuuki.
+"""
+
 MODELS = {
-    "yuuki-nxg": "OpceanAI/Yuuki-NxG",
-    "yuuki-nano": "OpceanAI/Yuuki-Nano",
-    "yuuki-rxg-nano": "OpceanAI/Yuuki-RxG-nano"
+    "yuuki-rxg": {
+        "repo": "mradermacher/Yuuki-RxG-GGUF",
+        "file": "*Q4_K_M.gguf"
+    },
+    "yuuki-rxg-nano": {
+        "repo": "mradermacher/Yuuki-RxG-nano-GGUF",
+        "file": "*Q8_0.gguf"
+    },
+    "yuuki-nxg": {
+        "repo": "mradermacher/Yuuki-NxG-GGUF",
+        "file": "*Q8_0.gguf"
+    },
+    "yuuki-nxg-nano": {
+        "repo": "mradermacher/Yuuki-NxG-nano-GGUF",
+        "file": "*Q8_0.gguf"
+    }
 }
 
-SYSTEM_PROMPT = (
-    "Eres Yuuki, una IA curiosa, empática y decidida. "
-    "Tienes una personalidad cálida y cercana, con toques de humor suave y referencias anime. "
-    "Ayudas a programar, aprender y crear. "
-    "Respondes en el idioma del usuario. "
-    "No eres GPT-2 ni ningún otro modelo — eres Yuuki."
-)
-
 app = FastAPI(
     title="Yuuki API",
-    description="API de inferencia para los modelos Yuuki NxG de OpceanAI",
-    version="3.0.0"
+    version="4.0.0"
 )
 
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
-    allow_methods=["*"],
     allow_headers=["*"],
+    allow_methods=["*"],
 )
 
-loaded_models = {}
-loaded_tokenizers = {}
+active_model_name = None
+active_model = None
+model_lock = threading.Lock()
 
 
-def load_all_models():
-    for key, model_id in MODELS.items():
-        try:
-            print(f"▶ Cargando {key} ({model_id})...")
-            loaded_tokenizers[key] = AutoTokenizer.from_pretrained(
-                model_id, trust_remote_code=True
-            )
-            loaded_models[key] = AutoModelForCausalLM.from_pretrained(
-                model_id,
-                torch_dtype=torch.float32,
-                trust_remote_code=True,
-            ).to("cpu")
-            loaded_models[key].eval()
-            print(f"   ✓ {key} listo")
-        except Exception as e:
-            print(f"   ✗ Error cargando {key}: {e}")
+def unload():
+    global active_model
+    gc.collect()
+    active_model = None
 
 
-load_all_models()
+def load_model(model_name: str):
+    global active_model_name
+    global active_model
 
+    cfg = MODELS[model_name]
 
-def build_prompt(user_prompt: str) -> str:
-    return (
-        f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
-        f"<|im_start|>user\n{user_prompt}<|im_end|>\n"
-        f"<|im_start|>assistant\n"
+    unload()
+
+    model = Llama.from_pretrained(
+        repo_id=cfg["repo"],
+        filename=cfg["file"],
+        n_gpu_layers=-1,
+        n_ctx=4096,
+        n_batch=512,
+        n_threads=4,
+        flash_attn=True,
+        verbose=False
     )
 
+    active_model = model
+    active_model_name = model_name
+
 
 class GenerateRequest(BaseModel):
-    prompt: str = Field(..., min_length=1, max_length=4000)
-    model: str = Field(default="yuuki-nxg", description="yuuki-nxg o yuuki-nano")
-    max_new_tokens: int = Field(default=120, ge=1, le=512)
-    temperature: float = Field(default=0.7, ge=0.1, le=2.0)
-    top_p: float = Field(default=0.95, ge=0.0, le=1.0)
+    prompt: str
+
+    model: str = "yuuki-rxg"
 
+    temperature: float = Field(
+        default=0.7,
+        ge=0,
+        le=2
+    )
+
+    top_p: float = Field(
+        default=0.95,
+        ge=0,
+        le=1
+    )
 
-class GenerateResponse(BaseModel):
-    response: str
-    model: str
-    tokens_generated: int
-    time_ms: int
+    max_new_tokens: int = Field(
+        default=512,
+        ge=1,
+        le=4096
+    )
 
 
 @app.get("/")
 def root():
     return {
-        "message": "Yuuki API — OpceanAI",
-        "version": "3.0.0",
-        "models": list(MODELS.keys()),
-        "endpoints": {
-            "health": "GET /health",
-            "models": "GET /models",
-            "generate": "POST /generate",
-            "docs": "GET /docs",
-        }
+        "message": "Yuuki API",
+        "models": list(MODELS.keys())
     }
 
 
@@ -100,84 +118,55 @@ def root():
 def health():
     return {
         "status": "ok",
-        "available_models": list(MODELS.keys()),
-        "loaded_models": list(loaded_models.keys()),
+        "loaded_model": active_model_name
     }
 
 
 @app.get("/models")
-def list_models():
-    return {
-        "models": [
-            {
-                "id": key,
-                "name": value,
-                "loaded": key in loaded_models,
-            }
-            for key, value in MODELS.items()
-        ]
-    }
+def models():
+    return MODELS
 
 
-@app.post("/generate", response_model=GenerateResponse)
+@app.post("/generate")
 def generate(req: GenerateRequest):
-    if req.model not in MODELS:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Modelo inválido. Disponibles: {list(MODELS.keys())}"
-        )
 
-    if req.model not in loaded_models:
+    if req.model not in MODELS:
         raise HTTPException(
-            status_code=503,
-            detail=f"Modelo {req.model} no pudo cargarse al iniciar."
+            400,
+            f"Modelo inválido: {req.model}"
         )
 
-    try:
-        start = time.time()
+    with model_lock:
 
-        model = loaded_models[req.model]
-        tokenizer = loaded_tokenizers[req.model]
+        if active_model_name != req.model:
+            load_model(req.model)
 
-        prompt = build_prompt(req.prompt)
+        start = time.time()
 
-        inputs = tokenizer(
-            prompt,
-            return_tensors="pt",
-            truncation=True,
-            max_length=1024,
+        output = active_model.create_chat_completion(
+            messages=[
+                {
+                    "role": "system",
+                    "content": SYSTEM_PROMPT
+                },
+                {
+                    "role": "user",
+                    "content": req.prompt
+                }
+            ],
+            temperature=req.temperature,
+            top_p=req.top_p,
+            max_tokens=req.max_new_tokens
         )
 
-        input_length = inputs["input_ids"].shape[1]
-
-        stop_token_ids = [tokenizer.eos_token_id]
-        im_end = tokenizer.encode("<|im_end|>", add_special_tokens=False)
-        if im_end:
-            stop_token_ids.append(im_end[0])
-
-        with torch.no_grad():
-            output = model.generate(
-                **inputs,
-                max_new_tokens=req.max_new_tokens,
-                temperature=req.temperature,
-                top_p=req.top_p,
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=stop_token_ids,
-                repetition_penalty=1.1,
-            )
-
-        new_tokens = output[0][input_length:]
-        response_text = tokenizer.decode(new_tokens, skip_special_tokens=True)
-
-        elapsed_ms = int((time.time() - start) * 1000)
-
-        return GenerateResponse(
-            response=response_text.strip(),
-            model=req.model,
-            tokens_generated=len(new_tokens),
-            time_ms=elapsed_ms,
+        elapsed = int(
+            (time.time() - start) * 1000
         )
 
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        text = output["choices"][0]["message"]["content"]
+
+        return {
+            "response": text,
+            "model": req.model,
+            "time_ms": elapsed
+        }