Add data transformer Gradio demo

README.md

@@ -1,13 +1,42 @@
 ---
-title: SLM  Transform
-emoji: 👁
-colorFrom: purple
-colorTo: yellow
+title: SLM 264M Data Transformer
+emoji: 🔄
+colorFrom: blue
+colorTo: indigo
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.12'
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SLM-264M Data Transformer — demo Space
+
+A Gradio front-end for the `slm-264m` data-transformer fine-tune: convert data between
+**CSV, JSON, and XML** with a 264M-parameter model trained from scratch, running on CPU.
+
+## How it's wired
+
+- **Code lives here** (`app.py`, `model.py`).
+- **Weights live in the model repo** and are downloaded at startup with
+  `huggingface_hub.hf_hub_download`. Set the repo in `app.py`:
+  ```python
+  MODEL_REPO = "uday210/slm-264m-transform"   # your fine-tune
+  CKPT_FILENAME = "slm-264m-transform.pt"
+  TOKENIZER_FILENAME = "tokenizer.json"
+  ```
+- If the model repo is **private**, add an `HF_TOKEN` secret under
+  *Settings → Variables and secrets*.
+
+## Files you must add
+
+1. **`model.py`** — copy it from your training repo. This Space cannot load the checkpoint
+   without your exact `GPT`/`GPTConfig` classes (custom GQA + RoPE architecture).
+   It must be import-safe (no training at file top level).
+2. Make sure the model repo actually contains the checkpoint and `tokenizer.json`.
+
+## Notes / limits
+
+- CPU inference: a few tok/s. Fine for a demo. For speed, upgrade the Space hardware or
+  wire in the KV-cache generation path from your deployment notes.
+- The prompt template in `app.py` must match how the fine-tune was trained.
+- Numeric transcription is the model's known weak spot; structure and text are reliable.

app.py

@@ -0,0 +1,139 @@
+"""
+Gradio demo for the SLM-264M data transformer fine-tune.
+
+Pattern: code lives here (in the Space), weights live in the model repo and are
+pulled at startup with huggingface_hub. You bundle model.py alongside this file.
+
+Files this Space expects:
+  - app.py            (this file)
+  - model.py          (YOUR architecture file — copy it from your training repo)
+  - requirements.txt
+  - README.md         (Space card, has the sdk: gradio frontmatter)
+
+Weights + tokenizer are downloaded from MODEL_REPO below. If that repo is private,
+add an HF_TOKEN secret to the Space (Settings -> Variables and secrets).
+"""
+
+import os
+import torch
+import gradio as gr
+from tokenizers import Tokenizer
+from huggingface_hub import hf_hub_download
+
+# ---- point these at your fine-tuned transform model repo -------------------
+# MODEL_REPO = os.environ.get("MODEL_REPO", "uday210/slm-264m-transform")  # <-- EDIT
+# CKPT_FILENAME = os.environ.get("CKPT_FILENAME", "slm-264m-transform.pt")  # <-- EDIT
+MODEL_REPO = "uday210/slm-transform-264m"   # <-- your actual transform model repo
+CKPT_FILENAME = "slm-transform-final.pt"     # <-- the real checkpoint filename
+TOKENIZER_FILENAME = os.environ.get("TOKENIZER_FILENAME", "tokenizer.json")
+
+# ---- prompt template: MUST match how the fine-tune was trained -------------
+# Your notes show the inference format as:
+#   "Convert this CSV to JSON:\n<data>\nOutput:\n"
+# If your training pairs used a different wording, change it here or the model
+# will be prompted off-distribution.
+PROMPT_TEMPLATE = "Convert this {src} to {tgt}:\n{data}\nOutput:\n"
+
+BLOCK_SIZE = 2048
+EOT_TOKEN = "<|endoftext|>"
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+if device == "cpu":
+    torch.set_num_threads(os.cpu_count() or 4)
+
+# ---------------------------------------------------------------------------
+# Load architecture + weights
+# ---------------------------------------------------------------------------
+from model import GPT  # noqa: E402  (your import-safe model.py)
+
+_token = os.environ.get("HF_TOKEN")  # only needed if the model repo is private
+
+ckpt_path = hf_hub_download(MODEL_REPO, CKPT_FILENAME, token=_token)
+tok_path = hf_hub_download(MODEL_REPO, TOKENIZER_FILENAME, token=_token)
+
+ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+model = GPT(ckpt["config"])
+model.load_state_dict(ckpt["model"])
+model.eval().to(device)
+
+tok = Tokenizer.from_file(tok_path)
+EOT_ID = tok.token_to_id(EOT_TOKEN)
+
+
+# ---------------------------------------------------------------------------
+# Generation (greedy by default — this task is deterministic transformation)
+# ---------------------------------------------------------------------------
+@torch.no_grad()
+def generate(prompt: str, max_new_tokens: int = 512, temperature: float = 0.0):
+    ids = tok.encode(prompt).ids
+    x = torch.tensor(ids, dtype=torch.long, device=device)[None]
+    start = len(ids)
+
+    for _ in range(max_new_tokens):
+        logits, _ = model(x[:, -BLOCK_SIZE:])
+        logits = logits[:, -1, :]
+        if temperature and temperature > 0:
+            probs = torch.softmax(logits / temperature, dim=-1)
+            nxt = torch.multinomial(probs, num_samples=1)
+        else:
+            nxt = logits.argmax(dim=-1, keepdim=True)
+        x = torch.cat([x, nxt], dim=1)
+        if EOT_ID is not None and nxt.item() == EOT_ID:
+            break
+        # stream the decoded suffix so far
+        yield tok.decode(x[0, start:].tolist())
+
+    yield tok.decode(x[0, start:].tolist())
+
+
+def transform(data, src_fmt, tgt_fmt, max_new_tokens, temperature):
+    data = (data or "").strip()
+    if not data:
+        yield "Paste some input data first."
+        return
+    if src_fmt == tgt_fmt:
+        yield "Source and target format are the same — pick different formats."
+        return
+    prompt = PROMPT_TEMPLATE.format(src=src_fmt, tgt=tgt_fmt, data=data)
+    for partial in generate(prompt, int(max_new_tokens), float(temperature)):
+        yield partial
+
+
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
+FORMATS = ["CSV", "JSON", "XML"]
+
+EXAMPLE_CSV = "name,city,age\nalice,london,31\nbob,paris,47\ncarol,tokyo,29"
+
+with gr.Blocks(title="SLM-264M Data Transformer", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        "# 🔄 SLM-264M Data Transformer\n"
+        "A 264M-parameter small language model, trained from scratch and fine-tuned to "
+        "convert between **CSV, JSON, and XML**. Runs on CPU — expect a few tokens/sec.\n\n"
+        "*Note: numbers can occasionally be miscopied; text and structure are reliable.*"
+    )
+    with gr.Row():
+        with gr.Column():
+            data_in = gr.Textbox(
+                label="Input data", value=EXAMPLE_CSV, lines=8,
+                placeholder="Paste CSV, JSON, or XML here…",
+            )
+            with gr.Row():
+                src = gr.Dropdown(FORMATS, value="CSV", label="From")
+                tgt = gr.Dropdown(FORMATS, value="JSON", label="To")
+            with gr.Accordion("Advanced", open=False):
+                max_tok = gr.Slider(64, 1024, value=512, step=64, label="Max new tokens")
+                temp = gr.Slider(
+                    0.0, 1.0, value=0.0, step=0.1,
+                    label="Temperature (0 = greedy, recommended for exact transforms)",
+                )
+            go = gr.Button("Transform", variant="primary")
+        with gr.Column():
+            out = gr.Textbox(label="Output", lines=14, show_copy_button=True)
+
+    go.click(transform, [data_in, src, tgt, max_tok, temp], out)
+    data_in.submit(transform, [data_in, src, tgt, max_tok, temp], out)
+
+if __name__ == "__main__":
+    demo.launch()

model.py

@@ -0,0 +1,186 @@
+"""
+model.py — Path B architecture for the ~250M SLM base.
+
+Modifications from build-nanogpt's GPT:
+  1. GQA (grouped-query attention) — fewer KV heads than query heads => small KV cache
+  2. RoPE (rotary position embeddings) — replaces learned wpe; extendable context
+  3. Tied embeddings — kept from baseline
+  4. Vocab size is a config field (set it to your trained tokenizer's size, e.g. 32768)
+
+Kept simple (LayerNorm + GELU) on purpose for a low-risk first real model.
+This file only DEFINES the model — it never trains on import, so it's safe to
+`from model import GPT, GPTConfig` from any script.
+"""
+import math
+from dataclasses import dataclass
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+# -----------------------------------------------------------------------------
+# RoPE helpers
+
+def build_rope_cache(seq_len, head_dim, device, base=10000.0):
+    """Precompute cos/sin tables for rotary embeddings. Shape: (seq_len, head_dim)."""
+    assert head_dim % 2 == 0, "head_dim must be even for RoPE"
+    theta = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    positions = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(positions, theta)                 # (seq_len, head_dim/2)
+    emb = torch.cat([freqs, freqs], dim=-1)               # (seq_len, head_dim)
+    return emb.cos(), emb.sin()
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rope(q, k, cos, sin):
+    # q, k: (B, n_head, T, head_dim); cos, sin: (T, head_dim)
+    cos = cos.unsqueeze(0).unsqueeze(0)   # (1, 1, T, head_dim)
+    sin = sin.unsqueeze(0).unsqueeze(0)
+    q_rot = (q * cos) + (rotate_half(q) * sin)
+    k_rot = (k * cos) + (rotate_half(k) * sin)
+    return q_rot, k_rot
+
+# -----------------------------------------------------------------------------
+# Grouped-Query Attention with RoPE
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        assert config.n_head % config.n_kv_head == 0, "n_head must be divisible by n_kv_head"
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = config.n_embd // config.n_head
+        self.n_rep = self.n_head // self.n_kv_head   # how many query heads share each KV head
+
+        # Q projects to full n_head; K and V project to only n_kv_head => smaller KV
+        self.q_proj = nn.Linear(config.n_embd, self.n_head * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
+        self.c_proj.NANOGPT_SCALE_INIT = 1
+
+    def forward(self, x, cos, sin):
+        B, T, C = x.size()
+        q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)     # (B, nh, T, hd)
+        k = self.k_proj(x).view(B, T, self.n_kv_head, self.head_dim).transpose(1, 2)  # (B, nkv, T, hd)
+        v = self.v_proj(x).view(B, T, self.n_kv_head, self.head_dim).transpose(1, 2)  # (B, nkv, T, hd)
+
+        # apply rotary embeddings to q and k
+        q, k = apply_rope(q, k, cos[:T], sin[:T])
+
+        # expand KV heads to match query heads (GQA): repeat each KV head n_rep times
+        k = k.repeat_interleave(self.n_rep, dim=1)   # (B, nh, T, hd)
+        v = v.repeat_interleave(self.n_rep, dim=1)   # (B, nh, T, hd)
+
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)  # flash attention
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+        return self.c_proj(y)
+
+class MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.gelu = nn.GELU(approximate='tanh')
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+        self.c_proj.NANOGPT_SCALE_INIT = 1
+
+    def forward(self, x):
+        return self.c_proj(self.gelu(self.c_fc(x)))
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(config.n_embd)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = nn.LayerNorm(config.n_embd)
+        self.mlp = MLP(config)
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.ln_1(x), cos, sin)
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+# -----------------------------------------------------------------------------
+
+@dataclass
+class GPTConfig:
+    block_size: int = 2048      # context length
+    vocab_size: int = 32768     # SET to your trained tokenizer's size (incl. FIM/special tokens)
+    n_layer: int = 24
+    n_head: int = 16            # query heads
+    n_kv_head: int = 4          # KV heads (GQA); n_head/n_kv_head = 4 query heads per KV head
+    n_embd: int = 1024
+
+class GPT(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.transformer = nn.ModuleDict(dict(
+            wte=nn.Embedding(config.vocab_size, config.n_embd),
+            h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            ln_f=nn.LayerNorm(config.n_embd),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        # tied embeddings
+        self.transformer.wte.weight = self.lm_head.weight
+
+        # RoPE cache (built lazily on first forward, cached on the module)
+        self.register_buffer("rope_cos", None, persistent=False)
+        self.register_buffer("rope_sin", None, persistent=False)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            std = 0.02
+            if hasattr(module, 'NANOGPT_SCALE_INIT'):
+                std *= (2 * self.config.n_layer) ** -0.5
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def _ensure_rope(self, T, device):
+        head_dim = self.config.n_embd // self.config.n_head
+        if self.rope_cos is None or self.rope_cos.size(0) < T or self.rope_cos.device != device:
+            cos, sin = build_rope_cache(max(T, self.config.block_size), head_dim, device)
+            self.rope_cos, self.rope_sin = cos, sin
+
+    def forward(self, idx, targets=None):
+        B, T = idx.size()
+        assert T <= self.config.block_size, f"sequence length {T} > block_size {self.config.block_size}"
+        self._ensure_rope(T, idx.device)
+        cos, sin = self.rope_cos.to(idx.device), self.rope_sin.to(idx.device)
+
+        x = self.transformer.wte(idx)     # (B, T, n_embd) — no positional embedding added; RoPE handles it
+        for block in self.transformer.h:
+            x = block(x, cos, sin)
+        x = self.transformer.ln_f(x)
+        logits = self.lm_head(x)
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+        return logits, loss
+
+    def num_params(self):
+        n = sum(p.numel() for p in self.parameters())
+        # subtract tied lm_head (shares wte weight) to avoid double counting
+        n -= self.lm_head.weight.numel()
+        return n
+
+# -----------------------------------------------------------------------------
+# quick self-test: prints param count so you can tune the config to ~250M
+if __name__ == "__main__":
+    cfg = GPTConfig()
+    model = GPT(cfg)
+    print(f"config: n_layer={cfg.n_layer}, n_embd={cfg.n_embd}, "
+          f"n_head={cfg.n_head}, n_kv_head={cfg.n_kv_head}, vocab={cfg.vocab_size}")
+    print(f"total parameters: {model.num_params()/1e6:.1f}M")
+    # tiny forward sanity check on CPU
+    x = torch.randint(0, cfg.vocab_size, (2, 128))
+    logits, _ = model(x)
+    print(f"forward OK — logits shape {tuple(logits.shape)}")

requirements.txt

@@ -0,0 +1,5 @@
+torch
+tokenizers
+huggingface_hub
+gradio>=4.44.0
+numpy