Update app.py

app.py

@@ -5,17 +5,33 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import time
 
-# Definir todos los modelos disponibles
+# Modelos disponibles
 MODELS = {
+    # Serie NxG (actual)
+    "yuuki-nxg": "OpceanAI/Yuuki-NxG",
+    "yuuki-nano": "OpceanAI/Yuuki-Nano",
+    # Serie Pre-NxG (legado)
     "yuuki-best": "OpceanAI/Yuuki-best",
     "yuuki-3.7": "OpceanAI/Yuuki-3.7",
-    "yuuki-v0.1": "OpceanAI/Yuuki-v0.1"
+    "yuuki-v0.1": "OpceanAI/Yuuki-v0.1",
 }
 
+# System prompt de Yuuki
+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."
+)
+
+# Modelos que usan ChatML (NxG)
+CHATML_MODELS = {"yuuki-nxg", "yuuki-nano"}
+
 app = FastAPI(
     title="Yuuki API",
-    description="Local inference API for Yuuki models",
-    version="1.0.0"
+    description="API de inferencia para los modelos Yuuki de OpceanAI",
+    version="2.0.0"
 )
 
 app.add_middleware(
@@ -25,31 +41,48 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Cache de modelos cargados
+# Cache de modelos
 loaded_models = {}
 loaded_tokenizers = {}
 
 
-def load_model(model_key: str):
-    """Lazy load: solo carga el modelo cuando se necesita"""
-    if model_key not in loaded_models:
-        print(f"Loading {model_key}...")
-        model_id = MODELS[model_key]
-        
-        loaded_tokenizers[model_key] = AutoTokenizer.from_pretrained(model_id)
-        loaded_models[model_key] = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            torch_dtype=torch.float32
-        ).to("cpu")
-        loaded_models[model_key].eval()
-        print(f"{model_key} ready!")
-    
-    return loaded_models[model_key], loaded_tokenizers[model_key]
+def load_all_models():
+    """Carga todos los modelos al iniciar"""
+    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}")
+
+
+# Cargar todos al arrancar
+load_all_models()
+
+
+def build_prompt(model_key: str, user_prompt: str) -> str:
+    """Construye el prompt según la serie del modelo"""
+    if model_key in CHATML_MODELS:
+        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"
+        )
+    return user_prompt  # Pre-NxG: prompt directo
 
 
 class GenerateRequest(BaseModel):
     prompt: str = Field(..., min_length=1, max_length=4000)
-    model: str = Field(default="yuuki-best", description="Model to use")
+    model: str = Field(default="yuuki-nxg", description="Modelo a usar")
     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)
@@ -65,13 +98,17 @@ class GenerateResponse(BaseModel):
 @app.get("/")
 def root():
     return {
-        "message": "Yuuki Local Inference API",
-        "models": list(MODELS.keys()),
+        "message": "Yuuki API — OpceanAI",
+        "version": "2.0.0",
+        "models": {
+            "nxg": [k for k in MODELS if k in CHATML_MODELS],
+            "legacy": [k for k in MODELS if k not in CHATML_MODELS],
+        },
         "endpoints": {
             "health": "GET /health",
             "models": "GET /models",
             "generate": "POST /generate",
-            "docs": "GET /docs"
+            "docs": "GET /docs",
         }
     }
 
@@ -81,7 +118,7 @@ def health():
     return {
         "status": "ok",
         "available_models": list(MODELS.keys()),
-        "loaded_models": list(loaded_models.keys())
+        "loaded_models": list(loaded_models.keys()),
     }
 
 
@@ -89,7 +126,12 @@ def health():
 def list_models():
     return {
         "models": [
-            {"id": key, "name": value} 
+            {
+                "id": key,
+                "name": value,
+                "series": "nxg" if key in CHATML_MODELS else "legacy",
+                "loaded": key in loaded_models,
+            }
             for key, value in MODELS.items()
         ]
     }
@@ -97,28 +139,42 @@ def list_models():
 
 @app.post("/generate", response_model=GenerateResponse)
 def generate(req: GenerateRequest):
-    # Validar que el modelo existe
     if req.model not in MODELS:
         raise HTTPException(
             status_code=400,
-            detail=f"Invalid model. Available: {list(MODELS.keys())}"
+            detail=f"Modelo inválido. Disponibles: {list(MODELS.keys())}"
         )
-    
+
+    if req.model not in loaded_models:
+        raise HTTPException(
+            status_code=503,
+            detail=f"Modelo {req.model} no pudo cargarse al iniciar."
+        )
+
     try:
         start = time.time()
 
-        # Cargar modelo (lazy load)
-        model, tokenizer = load_model(req.model)
+        model = loaded_models[req.model]
+        tokenizer = loaded_tokenizers[req.model]
+
+        prompt = build_prompt(req.model, req.prompt)
 
         inputs = tokenizer(
-            req.prompt,
+            prompt,
             return_tensors="pt",
             truncation=True,
-            max_length=1024
+            max_length=1024,
         )
 
         input_length = inputs["input_ids"].shape[1]
 
+        # Stop en <|im_end|> para modelos NxG
+        stop_token_ids = [tokenizer.eos_token_id]
+        if req.model in CHATML_MODELS:
+            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,
@@ -127,6 +183,7 @@ def generate(req: GenerateRequest):
                 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,
             )
 
@@ -139,8 +196,9 @@ def generate(req: GenerateRequest):
             response=response_text.strip(),
             model=req.model,
             tokens_generated=len(new_tokens),
-            time_ms=elapsed_ms
+            time_ms=elapsed_ms,
         )
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
+