Update HF-compatible weights

README.md

@@ -0,0 +1,49 @@
+---
+language: en
+tags:
+  - function-calling
+  - mobile-actions
+  - nanomind
+  - tool-use
+license: apache-2.0
+---
+
+# NanoMind — Mobile Actions (Function Calling)
+
+Fine-tuned on [google/mobile-actions](https://huggingface.co/datasets/google/mobile-actions)
+for on-device function calling (calendar, email, contacts, maps, flashlight, Wi-Fi).
+
+## Quick Start
+
+```python
+from hf_wrapper import NanoMindForFunctionCalling
+import tiktoken, torch, json
+
+model = NanoMindForFunctionCalling.from_pretrained(
+    "shawneil/NanoMind-MobileActions"
+)
+model.eval()
+enc = tiktoken.get_encoding("gpt2")
+
+# Build a prompt in the same format used during training
+tools_json = json.dumps([{"function": {"name": "create_calendar_event",
+    "description": "Creates a calendar event.",
+    "parameters": {"properties": {"title": {"type": "STRING"},
+    "datetime": {"type": "STRING"}}, "required": ["title","datetime"]}}}])
+
+prompt = (
+    "<|tools|>\n" + tools_json + "\n<|endtools|>\n"
+    "<|system|>\nYou are a function-calling assistant.\n<|endoftext|>\n"
+    "<|user|>\nSet a meeting tomorrow at 3 PM called 'Budget Review'.\n<|endoftext|>\n"
+    "<|assistant|>\n"
+)
+ids = torch.tensor([enc.encode_ordinary(prompt)], dtype=torch.long)
+out = model.generate_text(ids, max_new_tokens=80, temperature=0.0, top_k=1)
+print(enc.decode(out[0].tolist()))
+```
+
+## Training details
+- Base: NanoMind SFT-Dolly (shawneil/NanoMind-SFT-Dolly)
+- Dataset: google/mobile-actions (~9.65k rows)
+- Loss only on assistant/tool-call tokens
+- fp16 · torch.compile · 2× T4 DDP

config.json

@@ -0,0 +1,19 @@
+{
+  "model_type": "nanomind",
+  "architectures": [
+    "NanoMindForFunctionCalling"
+  ],
+  "task": "function_calling",
+  "train_step": 300,
+  "vocab_size": 50257,
+  "d_model": 512,
+  "n_heads": 8,
+  "n_kv_heads": 2,
+  "n_layers": 8,
+  "max_seq_len": 1024,
+  "ff_mult": 4,
+  "dropout": 0.0,
+  "use_moe": false,
+  "num_experts": 4,
+  "top_k_experts": 2
+}

hf_wrapper.py

@@ -0,0 +1,228 @@
+
+"""
+NanoMindForFunctionCalling — HuggingFace PreTrainedModel wrapper.
+
+Usage after downloading from HF hub:
+    from hf_wrapper import NanoMindForFunctionCalling
+    model = NanoMindForFunctionCalling.from_pretrained("shawneil/NanoMind-MobileActions")
+    # model is ready for inference, no separate architecture file needed
+"""
+
+import json, math, torch, torch.nn as nn, torch.nn.functional as F
+from dataclasses import dataclass, asdict
+from pathlib import Path
+
+try:
+    from transformers import PreTrainedModel, PretrainedConfig
+    HF_AVAILABLE = True
+except ImportError:
+    HF_AVAILABLE = False
+
+
+# ── Minimal standalone architecture ──────────────────────────────────
+@dataclass
+class ModelConfig:
+    vocab_size:    int   = 50257
+    d_model:       int   = 512
+    n_heads:       int   = 8
+    n_kv_heads:    int   = 2
+    n_layers:      int   = 8
+    max_seq_len:   int   = 512
+    ff_mult:       int   = 4
+    dropout:       float = 0.0
+    use_moe:       bool  = False
+    num_experts:   int   = 4
+    top_k_experts: int   = 2
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__(); self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+    def forward(self, x):
+        x32 = x.float()
+        return (x32 * x32.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
+                ).to(x.dtype).clone() * self.weight
+
+def _freqs_cis(head_dim, max_len, theta=10000.0):
+    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
+    t = torch.arange(max_len, device=freqs.device)
+    freqs = torch.outer(t, freqs)
+    return torch.polar(torch.ones_like(freqs), freqs)
+
+def _rope(xq, xk, fc):
+    def rot(x, f):
+        xc = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        return torch.view_as_real(xc * f[:x.shape[1]].unsqueeze(0).unsqueeze(2)
+                                   ).flatten(3).to(x.dtype)
+    return rot(xq, fc), rot(xk, fc)
+
+class GQAttn(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.nh = cfg.n_heads; self.nkv = cfg.n_kv_heads
+        self.hd = cfg.d_model // cfg.n_heads
+        self.q  = nn.Linear(cfg.d_model, cfg.n_heads    * self.hd, bias=False)
+        self.k  = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.hd, bias=False)
+        self.v  = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.hd, bias=False)
+        self.o  = nn.Linear(cfg.n_heads * self.hd, cfg.d_model,    bias=False)
+        self.drop = cfg.dropout
+    def forward(self, x, fc):
+        B, T, _ = x.shape
+        q = self.q(x).view(B, T, self.nh,  self.hd)
+        k = self.k(x).view(B, T, self.nkv, self.hd)
+        v = self.v(x).view(B, T, self.nkv, self.hd)
+        q, k = _rope(q, k, fc)
+        r = self.nh // self.nkv
+        k = k.repeat_interleave(r, 2); v = v.repeat_interleave(r, 2)
+        q, k, v = q.transpose(1,2), k.transpose(1,2), v.transpose(1,2)
+        out = F.scaled_dot_product_attention(q, k, v, None,
+              self.drop if self.training else 0., is_causal=True)
+        return self.o(out.transpose(1,2).contiguous().view(B, T, -1))
+
+class SwiGLU(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        h = (int(cfg.d_model * cfg.ff_mult * 2 / 3) + 63) // 64 * 64
+        self.w1 = nn.Linear(cfg.d_model, h, bias=False)
+        self.w2 = nn.Linear(h, cfg.d_model, bias=False)
+        self.w3 = nn.Linear(cfg.d_model, h, bias=False)
+    def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+class Block(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.an = RMSNorm(cfg.d_model); self.fn = RMSNorm(cfg.d_model)
+        self.attn = GQAttn(cfg); self.ff = SwiGLU(cfg)
+        self.drop = nn.Dropout(cfg.dropout)
+    def forward(self, x, fc):
+        x = x + self.drop(self.attn(self.an(x), fc))
+        return x + self.drop(self.ff(self.fn(x)))
+
+class _CoreModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop  = nn.Dropout(cfg.dropout)
+        self.blocks = nn.ModuleList([Block(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.embed.weight = self.lm_head.weight
+        self.register_buffer("freqs_cis",
+            _freqs_cis(cfg.d_model // cfg.n_heads, cfg.max_seq_len * 2))
+    def forward(self, idx, targets=None, loss_mask=None):
+        B, T = idx.shape
+        x = self.drop(self.embed(idx))
+        fc = self.freqs_cis[:T]
+        for blk in self.blocks: x = blk(x, fc)
+        logits = self.lm_head(self.norm(x))
+        loss = None
+        if targets is not None:
+            fl = logits.view(-1, logits.size(-1))
+            ft = targets.view(-1)
+            if loss_mask is not None:
+                m  = loss_mask.view(-1).bool()
+                fl = fl[m]; ft = ft[m]
+            loss = F.cross_entropy(fl, ft, ignore_index=-1)
+        return logits, loss
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens=200, temperature=0.8, top_k=50):
+        for _ in range(max_new_tokens):
+            ic = idx[:, -self.cfg.max_seq_len:]
+            logits, _ = self(ic)
+            logits = logits[:, -1, :] / temperature
+            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+            logits[logits < v[:, [-1]]] = float("-inf")
+            idx = torch.cat([idx, torch.multinomial(F.softmax(logits,-1), 1)], 1)
+        return idx
+
+
+# ── HF-compatible wrapper ─────────────────────────────────────────────
+if HF_AVAILABLE:
+    class NanoMindConfig(PretrainedConfig):
+        model_type = "nanomind"
+        def __init__(self, vocab_size=50257, d_model=512, n_heads=8,
+                     n_kv_heads=2, n_layers=8, max_seq_len=512,
+                     ff_mult=4, dropout=0.0, **kwargs):
+            super().__init__(**kwargs)
+            self.vocab_size  = vocab_size
+            self.d_model     = d_model
+            self.n_heads     = n_heads
+            self.n_kv_heads  = n_kv_heads
+            self.n_layers    = n_layers
+            self.max_seq_len = max_seq_len
+            self.ff_mult     = ff_mult
+            self.dropout     = dropout
+
+    class NanoMindForFunctionCalling(PreTrainedModel):
+        config_class = NanoMindConfig
+
+        def __init__(self, config: NanoMindConfig):
+            super().__init__(config)
+            cfg = ModelConfig(
+                vocab_size=config.vocab_size, d_model=config.d_model,
+                n_heads=config.n_heads, n_kv_heads=config.n_kv_heads,
+                n_layers=config.n_layers, max_seq_len=config.max_seq_len,
+                ff_mult=config.ff_mult, dropout=config.dropout,
+            )
+            self.model = _CoreModel(cfg)
+            self.post_init()
+
+        def forward(self, input_ids, labels=None, loss_mask=None, **kwargs):
+            logits, loss = self.model(input_ids, labels, loss_mask)
+            from transformers.modeling_outputs import CausalLMOutput
+            return CausalLMOutput(loss=loss, logits=logits)
+
+        @torch.no_grad()
+        def generate_text(self, idx, max_new_tokens=200, temperature=0.8, top_k=50):
+            return self.model.generate(idx, max_new_tokens, temperature, top_k)
+
+        @classmethod
+        def from_nanomind_checkpoint(cls, ckpt_path: str):
+            """Load from a raw NanoMind .pt checkpoint (no HF config needed)."""
+            ckpt   = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+            raw_cfg = ckpt.get("config", {})
+            hf_cfg = NanoMindConfig(
+                vocab_size  = raw_cfg.get("vocab_size",  50257),
+                d_model     = raw_cfg.get("d_model",     512),
+                n_heads     = raw_cfg.get("n_heads",     8),
+                n_kv_heads  = raw_cfg.get("n_kv_heads",  2),
+                n_layers    = raw_cfg.get("n_layers",    8),
+                max_seq_len = raw_cfg.get("max_seq_len", 512),
+                ff_mult     = raw_cfg.get("ff_mult",     4),
+                dropout     = raw_cfg.get("dropout",     0.0),
+            )
+            wrapper = cls(hf_cfg)
+            # remap keys: model.xxx → model.xxx (already correct)
+            state = ckpt["model"]
+            # if saved without wrapper prefix, add it
+            if not any(k.startswith("model.") for k in state):
+                state = {"model." + k: v for k, v in state.items()}
+            wrapper.load_state_dict(state, strict=True)
+            return wrapper
+
+else:
+    # Fallback when transformers not installed
+    class NanoMindForFunctionCalling(nn.Module):
+        def __init__(self, cfg: ModelConfig):
+            super().__init__()
+            self.model = _CoreModel(cfg)
+
+        @classmethod
+        def from_checkpoint(cls, ckpt_path):
+            ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+            raw_cfg = ckpt.get("config", {})
+            cfg = ModelConfig(**{k: v for k, v in raw_cfg.items()
+                                 if k in ModelConfig.__dataclass_fields__})
+            obj = cls(cfg)
+            state = ckpt["model"]
+            if not any(k.startswith("model.") for k in state):
+                state = {"model." + k: v for k, v in state.items()}
+            obj.load_state_dict(state, strict=True)
+            return obj
+
+        def forward(self, idx, targets=None, loss_mask=None):
+            return self.model(idx, targets, loss_mask)
+
+        def generate_text(self, idx, **kw):
+            return self.model.generate(idx, **kw)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:306542b1cd779ac09cb716c291fb89011f27a229cb68919e371a8a03009e6582
+size 296596728