Standardize TinyBrain bundle

README.md

@@ -1,20 +1,58 @@
 ---
 tags:
   - coreml
-  - Apple
+  - apple
   - watchos
-  - mlpackage
+  - tinybrain
 library_name: coreml
 ---
 
-# TinyGPT2 for CoreML
+# TinyGPT2 for TinyBrain
 
-This is a distilled GPT-2 model converted to CoreML as an ML Program (`.mlpackage`),
-compatible with Apple Watch (watchOS 10+) and other Apple devices.
+This repository packages `sshleifer/tiny-gpt2` for the TinyBrain runtime on Apple Watch.
 
-- Base model: `sshleifer/tiny-gpt2`
-- Format: `.mlpackage`
-- Input: tokenized text (input_ids)
-- Output: logits for language modeling
-- Compatible with `coremltools` and Apple Watch
+## TinyBrain bundle contract
 
+The conversion script produces a single bundle archive:
+
+- `tinygpt2_only_logits.mlmodelc.zip`
+
+The zip layout is:
+
+- `tinygpt2_only_logits.mlmodelc/`
+- `tokenizer/`
+- `metadata.json`
+
+## Runtime interface
+
+- Task: causal language modeling
+- Inputs:
+  - `input_ids` (`Int32`, shape `1 x 16`)
+  - `attention_mask` (`Int32`, shape `1 x 16`)
+- Output:
+  - `logits`
+
+## Files included in `tokenizer/`
+
+- `tokenizer.json` when provided by Transformers
+- `vocab.json`
+- `merges.txt`
+- `special_tokens_map.json`
+- `tokenizer_config.json`
+- `token_decoder.json`
+
+## Build
+
+Run:
+
+```bash
+python3 convert_tinygpt2_to_coreml.py
+```
+
+The script:
+
+1. Downloads the model and tokenizer from Hugging Face.
+2. Converts the model to Core ML.
+3. Compiles the package to `.mlmodelc`.
+4. Writes TinyBrain runtime metadata.
+5. Packages the compiled model, tokenizer, and metadata into one zip.

convert_tinygpt2_to_coreml.py

@@ -0,0 +1,119 @@
+import json
+import os
+import shutil
+import subprocess
+import zipfile
+
+import coremltools as ct
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+MODEL_NAME = "sshleifer/tiny-gpt2"
+BUNDLE_NAME = "tinygpt2_only_logits"
+SEQ_LEN = 16
+
+MLPACKAGE_NAME = f"{BUNDLE_NAME}.mlpackage"
+COMPILED_DIR = "CompiledModel"
+COMPILED_NAME = f"{BUNDLE_NAME}.mlmodelc"
+ZIP_NAME = f"{COMPILED_NAME}.zip"
+TOKENIZER_DIR = "tokenizer"
+METADATA_NAME = "metadata.json"
+
+
+def safe_rm(path: str) -> None:
+    if os.path.isdir(path):
+        shutil.rmtree(path, ignore_errors=True)
+    elif os.path.exists(path):
+        os.remove(path)
+
+
+def write_runtime_metadata(tokenizer, output_path: str) -> None:
+    metadata = {
+        "runtime": "tinybrain-causallm-v1",
+        "model_id": MODEL_NAME,
+        "bundle_name": BUNDLE_NAME,
+        "architecture": "gpt2",
+        "sequence_length": SEQ_LEN,
+        "vocab_size": len(tokenizer),
+        "input_names": ["input_ids", "attention_mask"],
+        "output_name": "logits",
+        "tokenizer_dir": "tokenizer",
+        "requires_compiled_model": True,
+    }
+    with open(output_path, "w", encoding="utf-8") as f:
+        json.dump(metadata, f, indent=2)
+
+
+for path in [ZIP_NAME, MLPACKAGE_NAME, COMPILED_DIR, TOKENIZER_DIR, METADATA_NAME]:
+    safe_rm(path)
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, use_safetensors=True)
+model.eval()
+
+
+class Wrapper(nn.Module):
+    def __init__(self, base_model):
+        super().__init__()
+        self.base_model = base_model
+
+    def forward(self, input_ids, attention_mask):
+        input_ids = input_ids[:, :SEQ_LEN]
+        attention_mask = attention_mask[:, :SEQ_LEN]
+        outputs = self.base_model(input_ids=input_ids, attention_mask=attention_mask)
+        return outputs.logits
+
+
+wrapped = Wrapper(model).eval()
+example = tokenizer(
+    "hello coreml",
+    return_tensors="pt",
+    padding="max_length",
+    truncation=True,
+    max_length=SEQ_LEN,
+)
+input_ids = example["input_ids"].to(dtype=torch.int32)
+attention_mask = example["attention_mask"].to(dtype=torch.int32)
+
+traced = torch.jit.trace(wrapped, (input_ids, attention_mask))
+mlmodel = ct.convert(
+    traced,
+    inputs=[
+        ct.TensorType(name="input_ids", shape=(1, SEQ_LEN), dtype=np.int32),
+        ct.TensorType(name="attention_mask", shape=(1, SEQ_LEN), dtype=np.int32),
+    ],
+    outputs=[ct.TensorType(name="logits")],
+    convert_to="mlprogram",
+    minimum_deployment_target=ct.target.iOS17,
+)
+mlmodel.save(MLPACKAGE_NAME)
+
+subprocess.run(["xcrun", "coremlc", "compile", MLPACKAGE_NAME, COMPILED_DIR], check=True)
+compiled_path = os.path.join(COMPILED_DIR, COMPILED_NAME)
+
+tokenizer.save_pretrained(TOKENIZER_DIR)
+with open(os.path.join(TOKENIZER_DIR, "token_decoder.json"), "w", encoding="utf-8") as f:
+    decoder = {str(token_id): token for token, token_id in tokenizer.get_vocab().items()}
+    json.dump(decoder, f, ensure_ascii=False)
+
+write_runtime_metadata(tokenizer, METADATA_NAME)
+
+with zipfile.ZipFile(ZIP_NAME, "w", zipfile.ZIP_DEFLATED) as archive:
+    for root, _, files in os.walk(compiled_path):
+        for filename in files:
+            full_path = os.path.join(root, filename)
+            rel_path = os.path.relpath(full_path, compiled_path)
+            archive.write(full_path, os.path.join(COMPILED_NAME, rel_path))
+    for root, _, files in os.walk(TOKENIZER_DIR):
+        for filename in files:
+            full_path = os.path.join(root, filename)
+            rel_path = os.path.relpath(full_path, TOKENIZER_DIR)
+            archive.write(full_path, os.path.join("tokenizer", rel_path))
+    archive.write(METADATA_NAME, METADATA_NAME)
+
+print(f"OK: {ZIP_NAME}")

metadata.json

@@ -0,0 +1,15 @@
+{
+  "runtime": "tinybrain-causallm-v1",
+  "model_id": "sshleifer/tiny-gpt2",
+  "bundle_name": "tinygpt2_only_logits",
+  "architecture": "gpt2",
+  "sequence_length": 16,
+  "vocab_size": 50257,
+  "input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "output_name": "logits",
+  "tokenizer_dir": "tokenizer",
+  "requires_compiled_model": true
+}

tinygpt2_only_logits.mlmodelc.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6552b89db61be95a22fb94961c5b404ade2011c1a5abeaabce383c741dbb61ef
+size 1773389

tokenizer/special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<|endoftext|>",
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer/tokenizer_config.json

@@ -0,0 +1,23 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}