import os import json import time import torch import torch.nn as nn import torch.nn.functional as F from fastapi import FastAPI, HTTPException from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse from pydantic import BaseModel from typing import Optional from tokenizers import Tokenizer from huggingface_hub import hf_hub_download, list_repo_files from safetensors.torch import load_file from sse_starlette.sse import EventSourceResponse # ── CONFIG ────────────────────────────────────────────────────────────────── HF_TOKEN = os.environ.get("HF_TOKEN", "") MODEL_REPO = "hugging-science/Nova-nano-2-chktps" DATA_REPO = "Bc-AI/nova1_data" DEVICE = "cuda" if torch.cuda.is_available() else "cpu" app = FastAPI(title="Nova-2-Nano Think Twice Demo", version="2.0.0") os.makedirs("static", exist_ok=True) app.mount("/static", StaticFiles(directory="static"), name="static") # ── MODEL ARCHITECTURE (Pure NTA) ─────────────────────────────────────────── class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.scale = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x): return F.rms_norm(x, self.scale.shape, self.scale, self.eps) def precompute_freqs_cis(head_dim, max_len, theta=10_000.0, device=None): freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim)) t = torch.arange(max_len, dtype=torch.float32, device=device) freqs = torch.outer(t, freqs) return torch.cos(freqs), torch.sin(freqs) def apply_rope(xq, xk, cos, sin): L = xq.shape[1] c = torch.cat([cos[:L], cos[:L]], -1).unsqueeze(0).unsqueeze(2) s = torch.cat([sin[:L], sin[:L]], -1).unsqueeze(0).unsqueeze(2) def rot(x): x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:] return torch.cat([-x2, x1], dim=-1) xq_f, xk_f = xq.float(), xk.float() return ((xq_f*c + rot(xq_f)*s).to(xq.dtype), (xk_f*c + rot(xk_f)*s).to(xk.dtype)) class InternalThinkGate(nn.Module): def __init__(self, d_model): super().__init__() self.probe = nn.Linear(d_model, 1, bias=False) self.blend_weight = nn.Parameter(torch.tensor(0.3)) def forward(self, x, sublayer_fn): out1 = sublayer_fn(x) conf = torch.sigmoid(self.probe(out1[:, -1, :])).mean() # Inference: only re-run if below threshold if conf < 0.80: out2 = sublayer_fn(x) alpha = torch.sigmoid(self.blend_weight) return alpha * out1 + (1 - alpha) * out2, conf.item() return out1, conf.item() class SwiGLU(nn.Module): def __init__(self, d_model, ffn_hidden): super().__init__() self.gate = nn.Linear(d_model, ffn_hidden, bias=False) self.up = nn.Linear(d_model, ffn_hidden, bias=False) self.down = nn.Linear(ffn_hidden, d_model, bias=False) def forward(self, x): return self.down(F.silu(self.gate(x)) * self.up(x)) class SlidingWindowAttention(nn.Module): def __init__(self, cfg): super().__init__() self.nh, self.hd = cfg['n_heads'], cfg['head_dim'] self.window = cfg['sliding_window'] D = self.nh * self.hd self.q = nn.Linear(cfg['d_model'], D, bias=False) self.k = nn.Linear(cfg['d_model'], D, bias=False) self.v = nn.Linear(cfg['d_model'], D, bias=False) self.o = nn.Linear(D, cfg['d_model'], bias=False) def forward(self, x, cos, sin): B, L, _ = x.shape q = self.q(x).view(B, L, self.nh, self.hd) k = self.k(x).view(B, L, self.nh, self.hd) v = self.v(x).view(B, L, self.nh, self.hd) q, k = apply_rope(q, k, cos, sin) mask = torch.ones(L, L, device=x.device, dtype=torch.bool) rows = torch.arange(L, device=x.device).unsqueeze(1) cols = torch.arange(L, device=x.device).unsqueeze(0) causal = cols <= rows windowed = (rows - cols) < self.window mask = ~(causal & windowed) q, k, v = q.transpose(1,2), k.transpose(1,2), v.transpose(1,2) out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) return self.o(out.transpose(1,2).contiguous().view(B, L, -1)) class LocalAttention(nn.Module): def __init__(self, cfg): super().__init__() self.nh, self.hd = cfg['n_heads'], cfg['head_dim'] self.radius = cfg['local_radius'] D = self.nh * self.hd self.q = nn.Linear(cfg['d_model'], D, bias=False) self.k = nn.Linear(cfg['d_model'], D, bias=False) self.v = nn.Linear(cfg['d_model'], D, bias=False) self.o = nn.Linear(D, cfg['d_model'], bias=False) def forward(self, x, cos, sin): B, L, _ = x.shape q = self.q(x).view(B, L, self.nh, self.hd) k = self.k(x).view(B, L, self.nh, self.hd) v = self.v(x).view(B, L, self.nh, self.hd) q, k = apply_rope(q, k, cos, sin) mask = torch.ones(L, L, device=x.device, dtype=torch.bool) rows = torch.arange(L, device=x.device).unsqueeze(1) cols = torch.arange(L, device=x.device).unsqueeze(0) causal = cols <= rows local_band = (rows - cols).abs() <= self.radius mask = ~(causal & local_band) q, k, v = q.transpose(1,2), k.transpose(1,2), v.transpose(1,2) out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) return self.o(out.transpose(1,2).contiguous().view(B, L, -1)) class GQA(nn.Module): def __init__(self, cfg): super().__init__() self.nh, self.nkv, self.hd = cfg['n_heads'], cfg['n_kv_heads'], cfg['head_dim'] self.ng = cfg['n_heads'] // cfg['n_kv_heads'] D, Dkv = cfg['n_heads']*cfg['head_dim'], cfg['n_kv_heads']*cfg['head_dim'] self.q = nn.Linear(cfg['d_model'], D, bias=False) self.k = nn.Linear(cfg['d_model'], Dkv, bias=False) self.v = nn.Linear(cfg['d_model'], Dkv, bias=False) self.o = nn.Linear(D, cfg['d_model'], bias=False) def forward(self, x, cos, sin): B, L, _ = x.shape q = self.q(x).view(B, L, self.nh, self.hd) k = self.k(x).view(B, L, self.nkv, self.hd) v = self.v(x).view(B, L, self.nkv, self.hd) q, k = apply_rope(q, k, cos, sin) q = q.transpose(1,2) k = k.transpose(1,2).repeat_interleave(self.ng, 1) v = v.transpose(1,2).repeat_interleave(self.ng, 1) out = F.scaled_dot_product_attention(q, k, v, is_causal=True) return self.o(out.transpose(1,2).contiguous().view(B, L, -1)) class NovaTripleAttention(nn.Module): def __init__(self, cfg): super().__init__() self.sliding = SlidingWindowAttention(cfg) self.gqa = GQA(cfg) self.local = LocalAttention(cfg) self.gate = nn.Linear(cfg['d_model'], 3, bias=False) self.proj = nn.Linear(cfg['d_model'], cfg['d_model'], bias=False) def forward(self, x, cos, sin): out_s = self.sliding(x, cos, sin) out_g = self.gqa(x, cos, sin) out_l = self.local(x, cos, sin) weights = F.softmax(self.gate(x), dim=-1) merged = (weights[:,:,0:1]*out_s + weights[:,:,1:2]*out_g + weights[:,:,2:3]*out_l) return self.proj(merged) class NovaNanoBlock(nn.Module): def __init__(self, cfg): super().__init__() self.attn = NovaTripleAttention(cfg) self.attn_norm = RMSNorm(cfg['d_model']) self.ffn_norm = RMSNorm(cfg['d_model']) self.ffn = SwiGLU(cfg['d_model'], cfg['ffn_hidden']) self.think_gate = InternalThinkGate(cfg['d_model']) def forward(self, x, cos, sin): attn_out, conf = self.think_gate( self.attn_norm(x), lambda h: self.attn(h, cos, sin) ) x = x + attn_out x = x + self.ffn(self.ffn_norm(x)) return x, conf class NovaNano(nn.Module): def __init__(self, cfg): super().__init__() self.cfg = cfg self.embed = nn.Embedding(cfg['vocab_size'], cfg['d_model']) self.layers = nn.ModuleList([NovaNanoBlock(cfg) for _ in range(cfg['n_layers'])]) self.norm = RMSNorm(cfg['d_model']) self.lm_head = nn.Linear(cfg['d_model'], cfg['vocab_size'], bias=False) self.lm_head.weight = self.embed.weight cos, sin = precompute_freqs_cis(cfg['head_dim'], cfg['max_len']) self.register_buffer("rope_cos", cos, persistent=False) self.register_buffer("rope_sin", sin, persistent=False) def forward(self, input_ids): B, L = input_ids.shape x = self.embed(input_ids) cos, sin = self.rope_cos[:L], self.rope_sin[:L] layer_confs = [] for layer in self.layers: x, conf = layer(x, cos, sin) layer_confs.append(conf) x = self.norm(x) logits = self.lm_head(x) avg_conf = sum(layer_confs) / len(layer_confs) return logits, avg_conf # ── GLOBALS ───────────────────────────────────────────────────────────────── model, tokenizer, CFG = None, None, None class CompletionRequest(BaseModel): prompt: str max_tokens: int = 150 temperature: float = 0.7 top_k: int = 50 think_twice: bool = True @app.on_event("startup") async def load_model(): global model, tokenizer, CFG print("🔧 Loading Nova-2-Nano...") tok_path = hf_hub_download(repo_id=DATA_REPO, filename="nova_tokenizer.json", repo_type="dataset", token=HF_TOKEN) tokenizer = Tokenizer.from_file(tok_path) meta_path = hf_hub_download(repo_id=DATA_REPO, filename="metadata.json", repo_type="dataset", token=HF_TOKEN) with open(meta_path) as f: meta = json.load(f) actual_vocab = meta.get('vocab_size', 50268) padded_vocab = ((actual_vocab + 63) // 64) * 64 CFG = { 'vocab_size': padded_vocab, 'd_model': 1024, 'n_heads': 8, 'n_kv_heads': 4, 'n_layers': 12, 'max_len': 2048, 'head_dim': 128, 'ffn_hidden': ((int(1024 * 8/3) + 63) // 64) * 64, 'sliding_window': 128, 'local_radius': 32, } files = list_repo_files(repo_id=MODEL_REPO, repo_type="model", token=HF_TOKEN) st_files = sorted([f for f in files if f.endswith('.safetensors')], reverse=True) pt_files = sorted([f for f in files if f.endswith('.model.pt') and 'best' in f], key=lambda x: int(x.split('_s')[-1].split('.')[0]) if '_s' in x else 0, reverse=True) ckpt_file = st_files[0] if st_files else pt_files[0] print(f"📦 Loading checkpoint: {ckpt_file}") ckpt_path = hf_hub_download(repo_id=MODEL_REPO, filename=ckpt_file, token=HF_TOKEN) if ckpt_file.endswith('.safetensors'): state_dict = load_file(ckpt_path) else: state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True) state_dict = {k.replace('_orig_mod.', ''): v for k, v in state_dict.items()} model = NovaNano(CFG).to(DEVICE).to(torch.bfloat16) model.load_state_dict(state_dict, strict=False) model.eval() print(f"✅ Nova-2-Nano loaded on {DEVICE}! ({sum(p.numel() for p in model.parameters())/1e9:.3f}B params)") @app.post("/v1/completions/stream") async def stream_completions(req: CompletionRequest): if model is None: raise HTTPException(503, "Model loading...") async def event_generator(): T_EOS = "<" + "|endoftext|" + ">" id_eos = tokenizer.token_to_id(T_EOS) input_ids = torch.tensor([tokenizer.encode(req.prompt).ids], dtype=torch.long, device=DEVICE) generated = input_ids.clone() passes_used = 0 start_time = time.perf_counter() last_avg_conf = 0.0 yield {"event": "metadata", "data": json.dumps({"model": "Nova-2-Nano", "think_twice": req.think_twice})} for _ in range(req.max_tokens): if generated.shape[1] > CFG['max_len']: generated = generated[:, -CFG['max_len']:] with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.bfloat16): logits, avg_conf = model(generated) last_avg_conf = avg_conf next_logits = logits[:, -1, :] / req.temperature # Repetition penalty for prev_id in generated[0][-15:]: if next_logits[0, prev_id] > 0: next_logits[0, prev_id] /= 1.2 else: next_logits[0, prev_id] *= 1.2 probs = F.softmax(next_logits, dim=-1) topk_p, topk_i = torch.topk(probs, req.top_k) next_id = topk_i[0, torch.multinomial(topk_p[0], 1)].item() # ThinkTwice visualization: report layer confidence if req.think_twice and avg_conf < 0.80: passes_used += 1 yield {"event": "thinking", "data": json.dumps({"avg_layer_conf": round(avg_conf, 3)})} await asyncio.sleep(0.01) generated = torch.cat([generated, torch.tensor([[next_id]], device=DEVICE)], dim=1) if next_id == id_eos: break clean_token = tokenizer.decode([next_id]).replace(T_EOS, "") yield {"event": "token", "data": json.dumps({"text": clean_token})} latency = (time.perf_counter() - start_time) * 1000 yield {"event": "done", "data": json.dumps({ "think_twice_triggers": passes_used, "final_confidence": round(last_avg_conf, 3), "latency_ms": round(latency, 2) })} return EventSourceResponse(event_generator()) @app.get("/", response_class=FileResponse) async def serve_frontend(): return FileResponse("static/index.html")