Upload app.py

app.py

@@ -7,7 +7,9 @@ Qué hace
 - `milo_core.py` debe estar junto a este app.py en el Space (el export del Colab 3 lo incluye). Si no está,
   se descarga del repo del modelo.
 - Usa la plantilla de chat con tokens especiales (<|system|> <|user|> <|assistant|> <|end|>), la MISMA del
-  fine-tuning, y el motor de generación de milo_core (KV-cache, min-p, no-repeat-ngram, etc.).
+  fine-tuning, y el motor de generación de milo_core (KV-cache, min-p, decodificación especulativa, etc.).
+- Si el modelo se entrenó con razonamiento (config.json: "thinking": true), PRIMERO piensa (<think>...</think>) y
+  DESPUÉS responde. El razonamiento no se muestra salvo que se pida (show_thinking) o MILO_SHOW_THINKING=1.
 - Endpoints: /chat, /chat/stream (SSE), /web_chat (búsqueda web + respuesta), /health, /v1/model_info,
   /web_search_debug. La UI de Gradio queda en "/".
 
@@ -49,7 +51,11 @@ MAX_NEW_TOKENS_HARD_CAP = int(os.environ.get("MILO_MAX_NEW_TOKENS_CAP", "400"))
 NUM_THREADS = int(os.environ.get("MILO_NUM_THREADS", str(os.cpu_count() or 2)))
 USE_CONTEXT = os.environ.get("MILO_USE_CONTEXT", "0") == "1"    # 1 = usa el historial de la conversación
 NUM_CANDIDATES = int(os.environ.get("MILO_NUM_CANDIDATES", "1"))  # >1 = best-of-N (en lote), N veces más costoso
-SYSTEM_PROMPT_ENV = os.environ.get("MILO_SYSTEM_PROMPT")         # "none" para desactivar; si no, el de config.json
+SYSTEM_PROMPT_ENV = os.environ.get("MILO_SYSTEM_PROMPT")
+THINK_DEFAULT = os.environ.get("MILO_THINK", "1") == "1"        # 1 = piensa antes de responder (si el modelo lo soporta)
+MAX_THINK_TOKENS = int(os.environ.get("MILO_MAX_THINK", "200"))
+SHOW_THINKING = os.environ.get("MILO_SHOW_THINKING", "0") == "1"
+SPEC_TOKENS = int(os.environ.get("MILO_SPEC", "4"))              # 0 = apaga la decodificación especulativa         # "none" para desactivar; si no, el de config.json
 
 TAVILY_API_KEY = os.environ.get("TAVILY_API_KEY", "").strip()
 SEARCH_TIMEOUT = float(os.environ.get("MILO_SEARCH_TIMEOUT", "8"))
@@ -77,6 +83,7 @@ _blocked: Optional[torch.Tensor] = None
 _load_error: Optional[str] = None
 _tok_info: dict = {}
 _stage: Optional[str] = None
+_tf: bool = False                # el modelo usa bloque <think>
 _GEN_LOCK = threading.Lock()   # una generación a la vez: en CPU, varias en paralelo solo se estorban
 
 
@@ -84,7 +91,7 @@ _GEN_LOCK = threading.Lock()   # una generación a la vez: en CPU, varias en par
 # Carga del modelo
 # ==========================================================================================
 def load_model():
-    global _model, _tok, _cfg, _system, _blocked, _load_error, _tok_info, _stage
+    global _model, _tok, _cfg, _system, _blocked, _load_error, _tok_info, _stage, _tf
     print(f"[Milo X1] Descargando modelo desde {MODEL_REPO_ID} (revision={MODEL_REVISION}) ...")
     try:
         from safetensors.torch import load_file
@@ -130,11 +137,12 @@ def load_model():
         _system = cfg_json.get("system_prompt", mc.DEFAULT_SYSTEM)
     _blocked = mc.make_blocked_ids(tok, cfg.vocab_size, "cpu", allow=tok.stop_ids)
     _model, _tok, _cfg, _stage = model, tok, cfg, cfg_json.get("stage")
+    _tf = bool(cfg_json.get("thinking", False))
     _tok_info = {"vocab_size": tok.vocab_size, "unicode_roundtrip": mc.check_tokenizer_unicode(tok, verbose=False)}
     _load_error = None
     print(f"[Milo X1] Modelo cargado (etapa: {_stage}). {model.num_params()/1e6:.1f}M parámetros | d_model={cfg.d_model} "
           f"| capas={cfg.n_layers} | heads={cfg.n_heads} (kv={cfg.n_kv_heads}) | contexto={cfg.context_length} "
-          f"| vocab={tok.vocab_size} | unicode/emoji={'OK' if _tok_info['unicode_roundtrip'] else 'CON PÉRDIDAS'}")
+          f"| pensamiento={'sí' if _tf else 'no'} | vocab={tok.vocab_size} | unicode/emoji={'OK' if _tok_info['unicode_roundtrip'] else 'CON PÉRDIDAS'}")
 
 
 # ==========================================================================================
@@ -185,31 +193,35 @@ def make_gen(params: dict, seed=None) -> mc.GenConfig:
         max_new_tokens=int(params["max_new_tokens"]), temperature=float(params["temperature"]),
         top_k=int(params["top_k"] or 0), top_p=float(params["top_p"] if params["top_p"] else 1.0),
         min_p=float(params.get("min_p", 0.0)), repetition_penalty=float(params["repetition_penalty"]),
-        no_repeat_ngram=int(params.get("no_repeat_ngram", 0)), seed=seed)
+        no_repeat_ngram=int(params.get("no_repeat_ngram", 0)), speculative=SPEC_TOKENS, seed=seed)
 
 
-def iter_text(ids: List[int], gen: mc.GenConfig, cancel: Optional[threading.Event] = None):
-    """Generador SÍNCRONO de fragmentos de texto en vivo. La generación corre en un hilo; si el consumidor
-    deja de leer (cliente desconectado), se cancela al siguiente token."""
+def iter_events(msgs: List[dict], gen: mc.GenConfig, think: bool, max_think: int, cancel: Optional[threading.Event] = None):
+    """Generador SÍNCRONO de eventos en vivo: ("think", texto) mientras razona y ("answer", texto) al responder.
+    La generación corre en un hilo; si el consumidor deja de leer (cliente desconectado) se cancela al siguiente token."""
     cancel = cancel or threading.Event()
     q: "queue.Queue" = queue.Queue()
-    dec = IncrementalDecoder(_tok)
+    d_think, d_ans = IncrementalDecoder(_tok), IncrementalDecoder(_tok)
 
-    def on_token(t):
-        if cancel.is_set():
-            raise _Cancel()
-        d = dec.push(t)
-        if d:
-            q.put(d)
+    def cb(kind, dec):
+        def on_token(t):
+            if cancel.is_set():
+                raise _Cancel()
+            d = dec.push(t)
+            if d:
+                q.put((kind, d))
+        return on_token
 
     def work():
         try:
             with _GEN_LOCK:
                 if not cancel.is_set():
-                    mc.generate(_model, [ids], gen, stop_ids=_tok.stop_ids, blocked_ids=_blocked, on_token=on_token)
-            tail = dec.flush()
-            if tail:
-                q.put(tail)
+                    mc.chat_generate(_model, _tok, msgs, system=_system, gen=gen, thinking_format=_tf, think=think,
+                                     max_think=max_think, on_think=cb("think", d_think), on_answer=cb("answer", d_ans))
+            for kind, dec in (("think", d_think), ("answer", d_ans)):
+                tail = dec.flush()
+                if tail:
+                    q.put((kind, tail))
         except _Cancel:
             pass
         except Exception as e:  # noqa
@@ -243,54 +255,51 @@ def score_candidate(text: str) -> float:
     return (len(set(bigrams)) / max(len(bigrams), 1)) * 0.5 + (len(set(words)) / len(words)) * 0.3 + min(len(words), 40) / 40 * 0.2
 
 
-def generate_text(ids: List[int], gen: mc.GenConfig, n: int = 1) -> str:
-    """Respuesta completa. Con n>1 genera n candidatas EN LOTE (una sola llamada) y elige la más coherente."""
-    n = max(1, n)
+def run_chat(msgs: List[dict], gen: mc.GenConfig, think: bool, max_think: int, n: int = 1) -> mc.ChatResult:
+    """Respuesta completa. Con n>1 y SIN razonamiento, genera n candidatas EN LOTE y elige la más coherente."""
+    active = _tf and think
     with _GEN_LOCK:
-        outs = mc.generate(_model, [ids] * n, gen, stop_ids=_tok.stop_ids, blocked_ids=_blocked)
-    texts = [_tok.decode(o).strip() for o in outs]
-    return max(texts, key=score_candidate) if n > 1 else texts[0]
+        if n > 1 and not active:
+            ids = _tok.encode_chat(msgs, system=_system, add_generation_prompt=True, thinking=_tf, think_prefix="closed")
+            blocked = mc.make_blocked_ids(_tok, _cfg.vocab_size, "cpu", allow=_tok.stop_ids)
+            outs = mc.generate(_model, [ids] * n, gen, stop_ids=_tok.stop_ids, blocked_ids=blocked)
+            texts = [_tok.decode(o).strip() for o in outs]
+            return mc.ChatResult("", max(texts, key=score_candidate))
+        return mc.chat_generate(_model, _tok, msgs, system=_system, gen=gen, thinking_format=_tf, think=think, max_think=max_think)
 
 
 # ==========================================================================================
-# Armado del prompt (plantilla de chat con tokens especiales)
+# Armado de mensajes (se ajustan al contexto; la plantilla de chat la aplica milo_core)
 # ==========================================================================================
-def _budget(max_new: int) -> int:
+def _reserve(max_new: int, think: bool, max_think: int) -> int:
     ctx = _cfg.context_length
-    return ctx - min(max_new, ctx // 2)      # mismo cálculo que mc.generate
+    r = min(max_new, ctx // 2)
+    if _tf and think:
+        r += min(max_think, ctx // 2 - 8)
+    return r
 
 
-def build_ids(history: List[dict], user_text: str, max_new: int) -> List[int]:
+def fit_msgs(history: List[dict], user_text: str, max_new: int, think: bool, max_think: int) -> List[dict]:
     """Si no cabe: descarta turnos viejos y luego recorta el INICIO del mensaje (se conserva el final)."""
-    budget = _budget(max_new)
-    hist, text = list(history), user_text
-    if len(text) > budget * 16:
-        text = text[-budget * 16:]
-    ids: List[int] = []
-    for _ in range(80):
-        ids = _tok.encode_chat(hist + [{"role": "user", "content": text}], system=_system, add_generation_prompt=True)
-        if len(ids) <= budget:
-            return ids
-        if hist:
-            hist = hist[2:]
-            continue
-        if len(text) <= 1:
-            break
-        text = text[-max(1, int(len(text) * 0.85)):]
-    return ids[:1] + ids[-(budget - 1):]
+    sess = mc.ChatSession(_model, _tok, system=_system, gen=mc.GenConfig(max_new_tokens=max_new),
+                          thinking_format=_tf, think=think, max_think=max_think)
+    sess.history = list(history)
+    if len(user_text) > 16000:
+        user_text = user_text[-16000:]
+    return sess._fit(user_text, think)
 
 
-def build_search_ids(question: str, entrada: str, max_new: int) -> List[int]:
+def fit_search_msgs(question: str, entrada: str, max_new: int, think: bool, max_think: int) -> List[dict]:
     """Pregunta + resultados web en un solo mensaje de usuario (igual que el SFT fusiona instruction + input).
     Si no cabe, se recortan los RESULTADOS, nunca la pregunta."""
-    budget = _budget(max_new)
-    ent, ids = entrada, []
+    budget = max(32, _cfg.context_length - _reserve(max_new, think, max_think) - 4)
+    ent = entrada
     for _ in range(60):
-        ids = _tok.encode_chat([{"role": "user", "content": f"{question}\n\n{ent}"}], system=_system, add_generation_prompt=True)
-        if len(ids) <= budget:
-            return ids
+        msgs = [{"role": "user", "content": f"{question}\n\n{ent}"}]
+        if len(_tok.encode_chat(msgs, system=_system, thinking=_tf, think_prefix="open")) <= budget:
+            return msgs
         ent = ent[: int(len(ent) * 0.9)].rstrip()
-    return ids[:1] + ids[-(budget - 1):]
+    return [{"role": "user", "content": question[-2000:]}]
 
 
 _LEGACY_RE = re.compile(r"###\s*Instrucci[oó]n:\s*\n(.*?)(?:\n\s*###\s*Entrada:\s*\n(.*?))?\s*\n\s*###\s*Respuesta:\s*$", re.S)
@@ -520,10 +529,6 @@ def rate_allow(key: str, limit: int = SEARCH_RATE_PER_MIN, window: float = 60.0)
 
 
 
-def build_search_user_text(question: str, entrada: str) -> str:
-    return f"{question}\n\n{entrada}"
-
-
 # ==========================================================================================
 # API FastAPI
 # ==========================================================================================
@@ -548,6 +553,9 @@ class ChatRequest(BaseModel):
     repetition_penalty: float = Field(CHAT_DEFAULTS["repetition_penalty"], ge=1.0, le=2.0)
     no_repeat_ngram: int = Field(0, ge=0, le=8)
     seed: Optional[int] = None
+    think: Optional[bool] = Field(None, description="Pensar antes de responder (por defecto MILO_THINK; solo si el modelo lo soporta)")
+    show_thinking: Optional[bool] = Field(None, description="Incluir el razonamiento en la respuesta (por defecto MILO_SHOW_THINKING)")
+    max_think_tokens: Optional[int] = Field(None, ge=8, le=1000)
 
 
 def _history_msgs(messages: List[ChatMessage]) -> Tuple[List[dict], str]:
@@ -566,10 +574,13 @@ def _history_msgs(messages: List[ChatMessage]) -> Tuple[List[dict], str]:
 
 
 def _resolve_request(req: ChatRequest):
-    """Devuelve (ids, GenConfig, search_mode)."""
+    """Devuelve (mensajes, GenConfig, search_mode, think, max_think, show_thinking)."""
+    given = req.model_fields_set
+    think = THINK_DEFAULT if req.think is None else bool(req.think)
+    max_think = req.max_think_tokens or MAX_THINK_TOKENS
+    show = SHOW_THINKING if req.show_thinking is None else bool(req.show_thinking)
     legacy = parse_legacy_prompt(req.prompt) or (
         parse_legacy_prompt(req.messages[-1].content) if (not req.prompt and req.messages and req.messages[-1].role == "user") else None)
-    given = req.model_fields_set
     if legacy is not None and legacy[1]:                      # con "Entrada" = modo búsqueda web
         question, entrada = legacy
         p = dict(SEARCH_DEFAULTS)
@@ -577,7 +588,7 @@ def _resolve_request(req: ChatRequest):
             if k in given:
                 p[k] = getattr(req, k)
         p["max_new_tokens"] = min(p["max_new_tokens"], MAX_NEW_TOKENS_HARD_CAP)
-        return build_search_ids(question, entrada, p["max_new_tokens"]), make_gen(p, req.seed), True
+        return fit_search_msgs(question, entrada, p["max_new_tokens"], think, max_think), make_gen(p, req.seed), True, think, max_think, show
     if legacy is not None:
         history, user_text = [], legacy[0]
     else:
@@ -586,7 +597,7 @@ def _resolve_request(req: ChatRequest):
         history, user_text = _history_msgs(req.messages)
     p = dict(max_new_tokens=req.max_new_tokens, temperature=req.temperature, top_k=req.top_k, top_p=req.top_p,
              min_p=req.min_p, repetition_penalty=req.repetition_penalty, no_repeat_ngram=req.no_repeat_ngram)
-    return build_ids(history, user_text, req.max_new_tokens), make_gen(p, req.seed), False
+    return fit_msgs(history, user_text, req.max_new_tokens, think, max_think), make_gen(p, req.seed), False, think, max_think, show
 
 
 @app.on_event("startup")
@@ -597,7 +608,8 @@ def _startup():
 @app.get("/health")
 def health():
     return {"status": "ok" if _model is not None else "model_not_loaded", "model_repo": MODEL_REPO_ID,
-            "quantized": QUANTIZE, "error": _load_error, "stage": _stage,
+            "quantized": QUANTIZE, "error": _load_error, "stage": _stage, "thinking_model": _tf,
+            "speculative": {"tokens": SPEC_TOKENS, **mc.SPEC_STATS},
             "web_search": {"tavily_key_set": bool(TAVILY_API_KEY), "fallback": "wikipedia"},
             "tokenizer": _tok_info or None}
 
@@ -610,6 +622,7 @@ def model_info():
             "d_model": _cfg.d_model, "n_layers": _cfg.n_layers, "n_heads": _cfg.n_heads, "n_kv_heads": _cfg.n_kv_heads,
             "context_length": _cfg.context_length, "vocab_size": _cfg.vocab_size, "quantized": QUANTIZE,
             "stage": _stage, "use_context": USE_CONTEXT, "num_candidates": NUM_CANDIDATES,
+            "thinking_model": _tf, "think_default": THINK_DEFAULT and _tf, "max_think_tokens": MAX_THINK_TOKENS,
             "system_prompt": _system, "tokenizer": _tok_info}
 
 
@@ -617,37 +630,47 @@ def model_info():
 def chat(req: ChatRequest):
     if _model is None:
         raise HTTPException(503, _load_error or "Modelo no cargado.")
-    ids, gen, search = _resolve_request(req)
+    msgs, gen, search, think, max_think, show = _resolve_request(req)
     t0 = time.time()
-    response = generate_text(ids, gen, n=1 if search else NUM_CANDIDATES)
-    return {"response": response, "elapsed_seconds": round(time.time() - t0, 2),
-            "mode": "web_search" if search else "chat"}
+    r = run_chat(msgs, gen, think, max_think, n=1 if search else NUM_CANDIDATES)
+    out = {"response": r.answer, "elapsed_seconds": round(time.time() - t0, 2), "mode": "web_search" if search else "chat"}
+    if show and _tf:
+        out["thinking"] = r.thinking
+    return out
 
 
 @app.post("/chat/stream")
 async def chat_stream(req: ChatRequest, request: Request):
     if _model is None:
         raise HTTPException(503, _load_error or "Modelo no cargado.")
-    ids, gen, search = _resolve_request(req)
+    msgs, gen, search, think, max_think, show = _resolve_request(req)
 
     async def event_stream():
         loop = asyncio.get_running_loop()
         cancel = threading.Event()
         try:
-            if search or NUM_CANDIDATES <= 1:
-                it = iter_text(ids, gen, cancel)
+            if search or NUM_CANDIDATES <= 1 or (_tf and think):
+                it = iter_events(msgs, gen, think, max_think, cancel)
                 sentinel = object()
+                announced = False
                 while True:
-                    piece = await loop.run_in_executor(None, next, it, sentinel)
-                    if piece is sentinel:
+                    ev = await loop.run_in_executor(None, next, it, sentinel)
+                    if ev is sentinel:
                         break
-                    yield f"data: {json.dumps({'delta': piece}, ensure_ascii=False)}\n\n"
+                    kind, piece = ev
+                    if kind == "answer":
+                        yield f"data: {json.dumps({'delta': piece}, ensure_ascii=False)}\n\n"
+                    elif show:
+                        yield f"data: {json.dumps({'thinking_delta': piece}, ensure_ascii=False)}\n\n"
+                    elif not announced:
+                        announced = True
+                        yield f"data: {json.dumps({'status': 'thinking'})}\n\n"
                     if await request.is_disconnected():
                         break
             else:
-                best = await loop.run_in_executor(None, generate_text, ids, gen, NUM_CANDIDATES)
-                for i in range(0, len(best), 4):
-                    yield f"data: {json.dumps({'delta': best[i:i + 4]}, ensure_ascii=False)}\n\n"
+                r = await loop.run_in_executor(None, run_chat, msgs, gen, think, max_think, NUM_CANDIDATES)
+                for i in range(0, len(r.answer), 4):
+                    yield f"data: {json.dumps({'delta': r.answer[i:i + 4]}, ensure_ascii=False)}\n\n"
                     await asyncio.sleep(0.01)
         except Exception as e:  # noqa
             yield f"data: {json.dumps({'error': str(e)}, ensure_ascii=False)}\n\n"
@@ -688,11 +711,12 @@ async def web_chat(req: WebChatRequest, request: Request):
     items = pick_results(raw, dedupe_domains=(provider != "wikipedia"))
     entrada, items = build_entrada(items)
     max_new = min(SEARCH_DEFAULTS["max_new_tokens"], MAX_NEW_TOKENS_HARD_CAP)
-    ids = build_search_ids(question, entrada, max_new)
+    max_think = min(MAX_THINK_TOKENS, 160)
+    msgs = fit_search_msgs(question, entrada, max_new, THINK_DEFAULT, max_think)
     gen = make_gen(dict(SEARCH_DEFAULTS, max_new_tokens=max_new))
     t0 = time.time()
     try:
-        reply = await loop.run_in_executor(None, generate_text, ids, gen, 1)
+        reply = (await loop.run_in_executor(None, run_chat, msgs, gen, THINK_DEFAULT, max_think, 1)).answer
     except Exception as e:  # noqa
         print(f"[Milo X1] web_chat: generación falló: {e}")
         return {"fallback": True, "notice": "No pude generar la respuesta con la búsqueda web ahora", "debug": str(e)[:300]}
@@ -744,20 +768,27 @@ def gradio_respond(message: str, history):
         yield _load_error or "Modelo no cargado. Revisa MODEL_REPO_ID del Space."
         return
     hist = _history_to_msgs(history)[-8:] if USE_CONTEXT else []
-    ids = build_ids(hist, message, CHAT_DEFAULTS["max_new_tokens"])
+    msgs = fit_msgs(hist, message, CHAT_DEFAULTS["max_new_tokens"], THINK_DEFAULT, MAX_THINK_TOKENS)
     gen = make_gen(CHAT_DEFAULTS)
-    if NUM_CANDIDATES > 1:
-        best = generate_text(ids, gen, n=NUM_CANDIDATES)
+    if NUM_CANDIDATES > 1 and not (_tf and THINK_DEFAULT):
+        best = run_chat(msgs, gen, False, MAX_THINK_TOKENS, n=NUM_CANDIDATES).answer
         for i in range(0, len(best), 4):
             yield best[: i + 4]
         if not best.strip():
             yield "(el modelo no generó texto — prueba otra pregunta o revisa el checkpoint)"
         return
-    partial = ""
-    for delta in iter_text(ids, gen):
-        partial += delta
-        yield partial
-    if not partial.strip():
+    thinking, answer = "", ""
+    if _tf and THINK_DEFAULT:
+        yield "💭 Pensando…"
+    for kind, piece in iter_events(msgs, gen, THINK_DEFAULT, MAX_THINK_TOKENS):
+        if kind == "think":
+            thinking += piece
+            if SHOW_THINKING:
+                yield f"💭 {thinking}"
+        else:
+            answer += piece
+            yield (f"💭 {thinking}\n\n{answer}" if (SHOW_THINKING and thinking) else answer)
+    if not answer.strip():
         yield "(el modelo no generó texto — prueba otra pregunta o revisa el checkpoint)"