app.py

# -*- coding: utf-8 -*-
"""Aether-7B-5Attn serving Space — FastAPI + static index.html, with live token streaming.

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, sys, threading, time
from fastapi import FastAPI
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")  # 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}
_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 load. Never raises to the request path."""
    if _state["model"] is not None or _state["status"] == "loading":
        return
    with _load_lock:
        if _state["model"] is not None:
            return
        _state["status"] = "loading"
        try:
            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"
            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

            # This checkpoint was instruction-tuned with attention_mask=None, so the model
            # never saw a mask during training. generate() builds one automatically, which
            # pushes it off-distribution and degenerates the output ("국가의 수도는 국가의
            # 수도입니다..."). Dropping the mask restores the trained behaviour. Measured
            # 2026-07-20: with mask -> degenerate, without -> correct on every probe.
            _fwd = model.forward

            def _forward_without_mask(*a, **kw):
                kw.pop("attention_mask", None)
                return _fwd(*a, **kw)

            model.forward = _forward_without_mask

            _state.update(model=model, tok=tok, device=dev, status="ready", error=None)
        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):
    prompt: str
    max_new_tokens: int | None = None
    temperature: float | None = 0.7


def _prep(req: "GenReq"):
    tok, model, dev = _state["tok"], _state["model"], _state["device"]
    prompt = (req.prompt or "").strip()[:4000]
    ids = None
    try:  # use the instruct chat template when available
        enc = tok.apply_chat_template([{"role": "user", "content": prompt}],
                                      add_generation_prompt=True, return_tensors="pt")
        # Transformers 5.x may hand back a BatchEncoding rather than a bare tensor. Passing that
        # into generate() dies on `inputs_tensor.shape[0]`, so always unwrap to the tensor.
        if hasattr(enc, "input_ids"):
            ids = enc.input_ids
        elif isinstance(enc, dict):
            ids = enc["input_ids"]
        else:
            ids = enc
    except Exception:
        ids = None
    if ids is None or not hasattr(ids, "shape"):
        ids = tok(prompt, return_tensors="pt").input_ids
    ids = ids.to(dev)  # batch_size = 1 ONLY
    # Greedy decoding is enforced. This checkpoint carries only light instruction
    # tuning, so sampling drifts off-distribution and hallucinates (e.g. inventing a
    # company profile when asked who it is). Deterministic decoding reproduces what
    # was actually trained. `temperature` is accepted for API compatibility but ignored.
    gen = dict(
        max_new_tokens=min(int(req.max_new_tokens or MAX_NEW), MAX_NEW),
        do_sample=False,
        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():
    import torch
    vram = None
    try:
        if torch.cuda.is_available():
            free, total = torch.cuda.mem_get_info()
            vram = {"free_gb": round(free / 1e9, 2), "total_gb": round(total / 1e9, 2),
                    "allocated_gb": round(torch.cuda.memory_allocated() / 1e9, 2)}
    except Exception:
        pass
    return {"status": _state["status"], "device": _state["device"], "model": MODEL_ID,
            "error": _state["error"], "gen_error": _state.get("gen_error"), "vram": vram}


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")

    from transformers import TextIteratorStreamer

    def run():
        # Acquire AND release inside the generator. Acquiring in the endpoint and releasing in
        # the generator's finally leaks the lock whenever the response is never consumed (client
        # disconnects before streaming starts) — which wedges the demo at "busy" forever.
        if not _gen_lock.acquire(timeout=2):
            yield "The demo is busy generating for another visitor — please retry in a moment."
            return
        try:
            tok, model, dev, ids, gen = _prep(req)
            streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True,
                                            timeout=120)
            th = threading.Thread(target=_safe_generate, daemon=True,
                                  args=(model, dict(input_ids=ids, streamer=streamer, **gen)))
            th.start()
            for chunk in streamer:
                yield chunk
            th.join(timeout=5)
        except Exception:
            yield "\n[generation stopped]"
        finally:
            _gen_lock.release()

    return StreamingResponse(run(), media_type="text/plain; charset=utf-8")


def _safe_generate(model, kwargs):
    """Never raise into the request path — but DO record why, or failures are invisible."""
    import torch, traceback
    try:
        torch.cuda.empty_cache()
    except Exception:
        pass
    try:
        with torch.no_grad():
            model.generate(**kwargs)
        _state["gen_error"] = None
    except Exception as e:
        # keep the full stack — a bare "AttributeError:" tells us nothing about where it came from
        _state["gen_error"] = ("%s: %s\n--- traceback ---\n%s"
                               % (type(e).__name__, str(e)[:200], traceback.format_exc()[-1500:]))
        traceback.print_exc()
    finally:
        try:
            torch.cuda.empty_cache()
        except Exception:
            pass


@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": _not_ready_message()})
    if (req.prompt or "").strip() == "":
        return JSONResponse(status_code=400, content={"ok": False, "message": "empty prompt"})
    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)
def index():
    try:
        return HTMLResponse(open("index.html", encoding="utf-8").read())
    except Exception:
        return HTMLResponse("<h1>Aether-7B-5Attn</h1><p>index.html missing.</p>")