# -*- 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("Aether-7B-5Attn
index.html missing.
")