Upload folder using huggingface_hub

README.md

@@ -1,12 +1,11 @@
 ---
-title: Nanochat Zerogpu Demo
-emoji: 🐢
-colorFrom: blue
-colorTo: gray
+title: nanochat ZeroGPU Demo
+emoji: 🤖
+colorFrom: indigo
+colorTo: blue
 sdk: gradio
-sdk_version: 5.49.1
+sdk_version: 4.41.0
 app_file: app.py
-pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+ZeroGPU demo Space serving nanochat SFT/MID/BASE models (CPU inference).

app.py

@@ -0,0 +1,78 @@
+
+import os
+import json
+import gradio as gr
+import torch
+from huggingface_hub import hf_hub_download
+from tokenizers import Tokenizer as HFTokenizer
+from gpt_infer import GPT, GPTConfig
+
+DEFAULT_MODEL = os.environ.get('NANOCHAT_DEFAULT_MODEL', 'loocorez/nanochat-sft-d20-step650')
+
+def load_model(repo_id: str):
+    cfg_path = hf_hub_download(repo_id, 'config.json')
+    with open(cfg_path, 'r') as f:
+        cfg = json.load(f)
+    model_config = GPTConfig(
+        sequence_len=cfg.get('n_ctx', 2048),
+        vocab_size=cfg['vocab_size'],
+        n_layer=cfg['n_layer'],
+        n_head=cfg['n_head'],
+        n_kv_head=cfg.get('n_kv_head', cfg['n_head']),
+        n_embd=cfg['n_embd'],
+    )
+    model = GPT(model_config)
+    model.eval()
+    try:
+        from safetensors.torch import load_file
+        weights_path = hf_hub_download(repo_id, 'model.safetensors')
+        sd = load_file(weights_path)
+    except Exception:
+        weights_path = hf_hub_download(repo_id, 'pytorch_model.bin')
+        sd = torch.load(weights_path, map_location='cpu')
+    model.load_state_dict(sd, strict=True, assign=True)
+    tok_path = hf_hub_download(repo_id, 'tokenizer.json')
+    tok = HFTokenizer.from_file(tok_path)
+    return model, tok
+
+model_cache = {}
+
+def get_model(repo_id: str):
+    if repo_id not in model_cache:
+        model_cache[repo_id] = load_model(repo_id)
+    return model_cache[repo_id]
+
+@torch.inference_mode()
+def generate(repo_id: str, prompt: str, max_tokens: int, temperature: float, top_k: int|None):
+    model, tok = get_model(repo_id)
+    bos_id = tok.token_to_id('<|bos|>')
+    ids = tok.encode(prompt).ids
+    if bos_id is not None:
+        ids = [bos_id] + ids
+    out_tokens = []
+    for token in model.generate(ids, max_tokens=max_tokens, temperature=temperature, top_k=top_k or None):
+        out_tokens.append(token)
+    text = tok.decode(out_tokens, skip_special_tokens=False)
+    return text
+
+with gr.Blocks() as demo:
+    gr.Markdown('# nanochat (ZeroGPU)')
+    gr.Markdown('Select a model and generate text.')
+    repo = gr.Dropdown(choices=[
+        'loocorez/nanochat-sft-d20-step650',
+        'loocorez/nanochat-mid-d20-step765',
+        'loocorez/nanochat-base-d20-step21400',
+    ], value=DEFAULT_MODEL, label='Model Repo')
+    prompt = gr.Textbox(label='Prompt', lines=6)
+    with gr.Row():
+        max_tokens = gr.Slider(1, 256, value=128, step=1, label='Max tokens')
+        temperature = gr.Slider(0.0, 1.5, value=0.8, step=0.05, label='Temperature')
+        top_k = gr.Slider(0, 100, value=40, step=1, label='Top-k (0=disabled)')
+    btn = gr.Button('Generate')
+    output = gr.Textbox(label='Output', lines=10)
+    btn.click(fn=lambda r,p,m,t,k: generate(r,p,int(m),float(t),int(k) if int(k)>0 else None),
+              inputs=[repo, prompt, max_tokens, temperature, top_k],
+              outputs=output)
+
+if __name__ == '__main__':
+    demo.launch()

gpt_infer.py

@@ -0,0 +1,190 @@
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass
+
+
+def norm(x: torch.Tensor) -> torch.Tensor:
+    return F.rms_norm(x, (x.size(-1),))
+
+
+def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    assert x.ndim == 4
+    d = x.shape[3] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    out = torch.cat([y1, y2], 3)
+    out = out.to(x.dtype)
+    return out
+
+
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    if n_rep == 1:
+        return x
+    bs, n_kv_heads, slen, head_dim = x.shape
+    return (
+        x[:, :, None, :, :]
+        .expand(bs, n_kv_heads, n_rep, slen, head_dim)
+        .reshape(bs, n_kv_heads * n_rep, slen, head_dim)
+    )
+
+
+@dataclass
+class GPTConfig:
+    sequence_len: int = 1024
+    vocab_size: int = 50304
+    n_layer: int = 12
+    n_head: int = 6
+    n_kv_head: int = 6
+    n_embd: int = 768
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config: GPTConfig, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = self.n_embd // self.n_head
+        assert self.n_embd % self.n_head == 0
+        assert self.n_kv_head <= self.n_head and self.n_head % self.n_kv_head == 0
+        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
+
+    def forward(self, x: torch.Tensor, cos_sin, kv_cache=None):
+        B, T, C = x.size()
+        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
+        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
+        cos, sin = cos_sin
+        q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)
+        q, k = norm(q), norm(k)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+        if kv_cache is not None:
+            k, v = kv_cache.insert_kv(self.layer_idx, k, v)
+        Tq = q.size(2)
+        Tk = k.size(2)
+        nrep = self.n_head // self.n_kv_head
+        k, v = repeat_kv(k, nrep), repeat_kv(v, nrep)
+        if kv_cache is None or Tq == Tk:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        elif Tq == 1:
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+        else:
+            attn_mask = torch.zeros((Tq, Tk), dtype=torch.bool, device=q.device)
+            prefix_len = Tk - Tq
+            if prefix_len > 0:
+                attn_mask[:, :prefix_len] = True
+            attn_mask[:, prefix_len:] = torch.tril(torch.ones((Tq, Tq), dtype=torch.bool, device=q.device))
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+        return y
+
+
+class MLP(nn.Module):
+    def __init__(self, config: GPTConfig):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, config: GPTConfig, layer_idx: int):
+        super().__init__()
+        self.attn = CausalSelfAttention(config, layer_idx)
+        self.mlp = MLP(config)
+
+    def forward(self, x: torch.Tensor, cos_sin, kv_cache=None) -> torch.Tensor:
+        x = x + self.attn(norm(x), cos_sin, kv_cache)
+        x = x + self.mlp(norm(x))
+        return x
+
+
+class GPT(nn.Module):
+    def __init__(self, config: GPTConfig):
+        super().__init__()
+        self.config = config
+        self.transformer = nn.ModuleDict({
+            'wte': nn.Embedding(config.vocab_size, config.n_embd),
+            'h': nn.ModuleList([Block(config, layer_idx) for layer_idx in range(config.n_layer)]),
+        })
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.rotary_seq_len = config.sequence_len * 10
+        head_dim = config.n_embd // config.n_head
+        cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
+        self.register_buffer('cos', cos, persistent=False)
+        self.register_buffer('sin', sin, persistent=False)
+        self.transformer.wte.to(dtype=torch.bfloat16)
+
+    def _precompute_rotary_embeddings(self, seq_len: int, head_dim: int, base: int = 10000, device=None):
+        if device is None:
+            device = self.transformer.wte.weight.device
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        t = torch.arange(seq_len, dtype=torch.float32, device=device)
+        freqs = torch.outer(t, inv_freq)
+        cos, sin = freqs.cos(), freqs.sin()
+        cos, sin = cos.bfloat16(), sin.bfloat16()
+        cos, sin = cos[None, :, None, :], sin[None, :, None, :]
+        return cos, sin
+
+    def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None, kv_cache=None, loss_reduction: str = 'mean'):
+        B, T = idx.size()
+        assert T <= self.cos.size(1)
+        assert idx.device == self.cos.device
+        assert self.cos.dtype == torch.bfloat16
+        T0 = 0 if kv_cache is None else kv_cache.get_pos()
+        cos_sin = self.cos[:, T0:T0+T], self.sin[:, T0:T0+T]
+        x = self.transformer.wte(idx)
+        x = norm(x)
+        for block in self.transformer.h:
+            x = block(x, cos_sin, kv_cache)
+        x = norm(x)
+        softcap = 15
+        if targets is not None:
+            logits = self.lm_head(x)
+            logits = softcap * torch.tanh(logits / softcap)
+            logits = logits.float()
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1, reduction=loss_reduction)
+            return loss
+        else:
+            logits = self.lm_head(x)
+            logits = softcap * torch.tanh(logits / softcap)
+            return logits
+
+    @torch.inference_mode()
+    def generate(self, tokens: list[int], max_tokens: int, temperature: float = 1.0, top_k: int | None = None, seed: int = 42):
+        assert isinstance(tokens, list)
+        device = self.transformer.wte.weight.device
+        rng = None
+        if temperature > 0:
+            rng = torch.Generator(device=device)
+            rng.manual_seed(seed)
+        ids = torch.tensor([tokens], dtype=torch.long, device=device)
+        for _ in range(max_tokens):
+            logits = self.forward(ids)
+            logits = logits[:, -1, :]
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            if temperature > 0:
+                logits = logits / max(temperature, 1e-6)
+                probs = F.softmax(logits, dim=-1)
+                next_ids = torch.multinomial(probs, num_samples=1, generator=rng)
+            else:
+                next_ids = torch.argmax(logits, dim=-1, keepdim=True)
+            ids = torch.cat((ids, next_ids), dim=1)
+            yield next_ids.item()

requirements.txt

@@ -0,0 +1,5 @@
+gradio==4.41.0
+huggingface_hub>=0.24.0
+tokenizers>=0.14.0
+safetensors>=0.4.3
+torch