app.py

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