serving: manual load + GenerationMixin + use_cache=False + streaming

app.py

@@ -1,48 +1,92 @@
 # -*- coding: utf-8 -*-
-"""Aether-7B-5Attn serving Space — FastAPI backend + static index.html.
+"""Aether-7B-5Attn serving Space — FastAPI + static index.html, with live token streaming.
 
-Serves the fully-open Aether model. The model is lazy-loaded on the first /api/generate call so
-the landing page (SEO/AEO) is instant. batch_size=1 is enforced (NSA branches ignore padding
-masks). Runs on GPU when available; on CPU it still works but is slow, so a clear status is
-returned. If the model can't be loaded (private + no token, or OOM), the page still serves and
-the API returns a friendly, honest message instead of crashing.
+Loading notes (this is a custom `aether_v2_7way` architecture, validated end-to-end on a T4):
+  * The repo ships the architecture in `aether_pkg/`; we snapshot it, import it, and subclass with
+    `GenerationMixin` (Transformers 5.x no longer gives `PreTrainedModel` a `.generate`).
+  * Random weight init is skipped (it is overwritten by the checkpoint) to keep cold-start sane.
+  * Weights load straight onto the GPU and the model moves over, so a 6.59B bf16 (~13.3 GB) model
+    fits a single T4 (16 GB) without a 2x memory spike.
+  * `use_cache=False`: the NSA attention branch uses custom KV-cache indices that stock
+    `DynamicCache` does not provide (fast KV-caching is a separate serving build), so generation
+    runs without a cache — correct but not fast. Output is streamed so tokens appear as produced;
+    `max_new_tokens` is capped and inference is batch_size = 1 (NSA ignores padding masks).
+
+The model pre-warms on startup. If it cannot load, the landing page still serves and the API
+returns an honest message instead of crashing.
 """
-import os, threading, time
+import os, sys, threading, time
 from fastapi import FastAPI
-from fastapi.responses import HTMLResponse, JSONResponse
+from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
 from pydantic import BaseModel
 
 MODEL_ID = os.environ.get("MODEL_ID", "FINAL-Bench/Aether-7B-5Attn-it")
-HF_TOKEN = os.environ.get("HF_TOKEN")  # Space secret; needed while the model repo is private
-MAX_NEW = int(os.environ.get("MAX_NEW_TOKENS", "256"))
+HF_TOKEN = os.environ.get("HF_TOKEN")  # optional now the repo is public
+MAX_NEW = min(int(os.environ.get("MAX_NEW_TOKENS", "64")), 200)  # no KV cache → keep it short
 
 app = FastAPI(title="Aether-7B-5Attn — Sovereign Open-Source AI")
 
 _state = {"model": None, "tok": None, "device": None, "status": "not_loaded", "error": None}
-_lock = threading.Lock()
+_load_lock = threading.Lock()
+_gen_lock = threading.Lock()   # single GPU demo → one generation at a time
+
+NOT_READY = ("Model is not available on this hardware yet. Serving Aether-7B-5Attn (6.59B) needs a "
+             "GPU. Details: ")
 
 
 def _load():
-    """Lazy, thread-safe model load. Sets _state; never raises to the request path."""
+    """Lazy, thread-safe load. Never raises to the request path."""
     if _state["model"] is not None or _state["status"] == "loading":
         return
-    with _lock:
+    with _load_lock:
         if _state["model"] is not None:
             return
         _state["status"] = "loading"
         try:
-            import torch
-            from transformers import AutoTokenizer, AutoModelForCausalLM
+            import torch, torch.nn as nn
+            from huggingface_hub import snapshot_download
+            from safetensors.torch import load_file
+            from transformers import AutoTokenizer, GenerationMixin
+            local = snapshot_download(MODEL_ID, token=HF_TOKEN)
+            if local not in sys.path:
+                sys.path.insert(0, local)
+            from aether_pkg.configuration_aether_v2_7way import AETHERV27wayConfig
+            from aether_pkg.modeling_aether_v2_7way import AETHERV27wayForCausalLM
+
+            class _AetherGen(AETHERV27wayForCausalLM, GenerationMixin):
+                pass
+
             dev = "cuda" if torch.cuda.is_available() else "cpu"
-            dtype = torch.bfloat16 if dev == "cuda" else torch.float32
-            tok = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN, trust_remote_code=True)
-            model = AutoModelForCausalLM.from_pretrained(
-                MODEL_ID, token=HF_TOKEN, trust_remote_code=True,
-                torch_dtype=dtype, low_cpu_mem_usage=True,
-            ).to(dev).eval()
+            tok = AutoTokenizer.from_pretrained(local, trust_remote_code=True)
+            cfg = AETHERV27wayConfig.from_pretrained(local)
+            cfg.use_cache = False
+            # skip pointless random init (overwritten by the checkpoint) → sane cold start
+            _sv = (nn.Linear.reset_parameters, nn.Embedding.reset_parameters)
+            nn.Linear.reset_parameters = lambda self: None
+            nn.Embedding.reset_parameters = lambda self: None
+            try:
+                torch.set_default_dtype(torch.bfloat16)
+                model = _AetherGen(cfg)
+            finally:
+                torch.set_default_dtype(torch.float32)
+                nn.Linear.reset_parameters, nn.Embedding.reset_parameters = _sv
+            sd = load_file(os.path.join(local, "model.safetensors"),
+                           device=(dev if dev == "cuda" else "cpu"))
+            model.load_state_dict(sd, strict=False)
+            del sd
+            if dev == "cuda":
+                torch.cuda.empty_cache()
+                model = model.to("cuda")
+            model.eval()
+            model.config.use_cache = False
             _state.update(model=model, tok=tok, device=dev, status="ready", error=None)
-        except Exception as e:  # honest degrade — page still works
-            _state.update(status="error", error=str(e)[:400])
+        except Exception as e:  # honest degrade — landing page still works
+            _state.update(status="error", error=str(e)[:500])
+
+
+@app.on_event("startup")
+def _prewarm():
+    threading.Thread(target=_load, daemon=True).start()
 
 
 class GenReq(BaseModel):
@@ -51,47 +95,97 @@ class GenReq(BaseModel):
     temperature: float | None = 0.7
 
 
+def _prep(req: "GenReq"):
+    tok, model, dev = _state["tok"], _state["model"], _state["device"]
+    prompt = (req.prompt or "").strip()[:4000]
+    try:  # use the instruct chat template when available
+        ids = tok.apply_chat_template([{"role": "user", "content": prompt}],
+                                      add_generation_prompt=True, return_tensors="pt").to(dev)
+    except Exception:
+        ids = tok(prompt, return_tensors="pt").input_ids.to(dev)  # batch_size = 1 ONLY
+    temp = float(req.temperature if req.temperature is not None else 0.7)
+    gen = dict(
+        max_new_tokens=min(int(req.max_new_tokens or MAX_NEW), MAX_NEW),
+        do_sample=temp > 0, temperature=max(temp, 0.01), top_p=0.9,
+        repetition_penalty=1.3, no_repeat_ngram_size=3, use_cache=False,
+        pad_token_id=getattr(tok, "eos_token_id", None),
+    )
+    return tok, model, dev, ids, gen
+
+
 @app.get("/api/health")
 def health():
     return {"status": _state["status"], "device": _state["device"], "model": MODEL_ID,
             "error": _state["error"]}
 
 
+def _not_ready_message():
+    return ("Model is still warming up — try again in a moment."
+            if _state["status"] in ("loading", "not_loaded") else NOT_READY + (_state["error"] or ""))
+
+
+@app.post("/api/stream")
+def stream(req: GenReq):
+    if _state["status"] in ("not_loaded", "loading"):
+        _load()
+    if _state["status"] != "ready":
+        return StreamingResponse(iter([_not_ready_message()]),
+                                 media_type="text/plain; charset=utf-8", status_code=503)
+    if (req.prompt or "").strip() == "":
+        return StreamingResponse(iter(["(empty prompt)"]), media_type="text/plain")
+    if not _gen_lock.acquire(blocking=False):
+        return StreamingResponse(
+            iter(["The demo is busy generating for another visitor — please retry in a moment."]),
+            media_type="text/plain; charset=utf-8", status_code=429)
+
+    from transformers import TextIteratorStreamer
+
+    def run():
+        try:
+            tok, model, dev, ids, gen = _prep(req)
+            streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True,
+                                            timeout=300)
+            th = threading.Thread(target=_safe_generate,
+                                  args=(model, dict(input_ids=ids, streamer=streamer, **gen)))
+            th.start()
+            for chunk in streamer:
+                yield chunk
+            th.join(timeout=5)
+        finally:
+            _gen_lock.release()
+
+    return StreamingResponse(run(), media_type="text/plain; charset=utf-8")
+
+
+def _safe_generate(model, kwargs):
+    import torch
+    try:
+        with torch.no_grad():
+            model.generate(**kwargs)
+    except Exception:
+        pass  # the streamer ends; the client sees whatever was produced
+
+
 @app.post("/api/generate")
 def generate(req: GenReq):
     if _state["status"] in ("not_loaded", "loading"):
         _load()
     if _state["status"] != "ready":
-        return JSONResponse(status_code=503, content={
-            "ok": False,
-            "status": _state["status"],
-            "message": ("Model is still loading — try again in a moment."
-                        if _state["status"] == "loading" else
-                        "Model is not available on this hardware yet. "
-                        "Serving Aether-7B-5Attn (6.59B) needs a GPU, and the repo must be public "
-                        "or a HF_TOKEN secret must be set. Details: " + (_state["error"] or "")),
-        })
-    import torch
-    tok, model, dev = _state["tok"], _state["model"], _state["device"]
-    prompt = (req.prompt or "").strip()[:4000]
-    if not prompt:
+        return JSONResponse(status_code=503, content={"ok": False, "status": _state["status"],
+                                                       "message": _not_ready_message()})
+    if (req.prompt or "").strip() == "":
         return JSONResponse(status_code=400, content={"ok": False, "message": "empty prompt"})
-    # batch_size=1 ONLY — NSA branches ignore padding masks.
-    ids = tok(prompt, return_tensors="pt").input_ids.to(dev)
-    t0 = time.time()
-    with torch.no_grad():
-        out = model.generate(
-            ids, max_new_tokens=min(req.max_new_tokens or MAX_NEW, 512),
-            do_sample=(req.temperature or 0) > 0, temperature=max(req.temperature or 0.7, 0.01),
-            top_p=0.9, repetition_penalty=1.1,
-            pad_token_id=getattr(tok, "eos_token_id", None),
-        )
-    text = tok.decode(out[0][ids.shape[1]:], skip_special_tokens=True)
-    dt = time.time() - t0
-    ntok = out.shape[1] - ids.shape[1]
-    return {"ok": True, "completion": text, "tokens": int(ntok),
-            "seconds": round(dt, 2), "tok_per_s": round(ntok / dt, 1) if dt else None,
-            "device": dev}
+    import torch
+    with _gen_lock:
+        tok, model, dev, ids, gen = _prep(req)
+        t0 = time.time()
+        with torch.no_grad():
+            out = model.generate(input_ids=ids, **gen)
+        text = tok.decode(out[0][ids.shape[1]:], skip_special_tokens=True)
+        dt = time.time() - t0
+        ntok = int(out.shape[1] - ids.shape[1])
+    return {"ok": True, "completion": text, "tokens": ntok, "seconds": round(dt, 2),
+            "tok_per_s": round(ntok / dt, 1) if dt else None, "device": dev}
 
 
 @app.get("/", response_class=HTMLResponse)