Upload folder using huggingface_hub

README.md

@@ -0,0 +1,92 @@
+---
+library_name: pytorch
+pipeline_tag: text-generation
+language:
+- en
+datasets:
+- roneneldan/TinyStories
+tags:
+- safetensors
+- custom-code
+---
+
+# Tiny CED
+
+Tiny CED is a small custom PyTorch encoder-decoder language model trained from
+scratch on the English [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories)
+dataset. This package contains inference-only weights in Safetensors format.
+
+This is not a Transformers `PreTrainedModel`. Load it with the included
+`model.py` and `safetensors.torch.load_model` as shown below.
+
+## Model details
+
+| Item | Value |
+| --- | ---: |
+| Independent parameters | 19,667,712 |
+| Vocabulary size | 8,192 |
+| Hidden size | 384 |
+| Encoder layers | 4 |
+| Decoder layers | 4 |
+| Attention heads | 6 |
+| Feed-forward size | 1,024 |
+| Local attention window | 64 |
+| Trained context length | 512 |
+| Weight dtype | FP32 |
+
+The token embedding and output head are tied. `model.safetensors` was exported
+with `safetensors.torch.save_model` so the shared tensor is stored only once.
+
+## Usage
+
+Install the runtime dependencies:
+
+```bash
+pip install -r requirements.txt
+```
+
+Run generation on CPU:
+
+```bash
+python generate.py \
+  --device cpu \
+  --text "Once upon a time, a little rabbit" \
+  --max-new-tokens 100
+```
+
+Use `--device cuda` when CUDA is available. Set `--temperature 0` for greedy
+decoding.
+
+## Training and evaluation
+
+The exported weights come from the existing `best.pt` checkpoint; no retraining
+was performed during export.
+
+- Optimizer: AdamW
+- Peak learning rate: 3e-4
+- Best-checkpoint training tokens: 40,004,782
+- Best-checkpoint step: 5,856
+- Validation loss: 1.896527
+- Validation perplexity: 6.662715
+
+Validation used the TinyStories validation split and the same 8,192-token BPE
+tokenizer included in this repository.
+
+## Files
+
+- `model.safetensors`: inference weights
+- `model.py`: exact custom PyTorch architecture
+- `config.json`: architecture and token IDs
+- `tokenizer.json`: BPE tokenizer
+- `generate.py`: minimal generation example
+
+## Limitations
+
+This is a small research model trained only on synthetic English children's
+stories. It may generate incorrect, repetitive, biased, unsafe, or otherwise
+inappropriate text. It is not suitable for factual, safety-critical, or
+production use without further evaluation.
+
+The source project did not specify a license for its code or weights. Choose an
+appropriate license and confirm the training-data terms before publishing this
+package publicly. TinyStories is listed as CDLA-Sharing-1.0 on its dataset page.

config.json

@@ -0,0 +1,18 @@
+{
+  "architectures": [
+    "CED"
+  ],
+  "model_type": "tiny-ced",
+  "vocab_size": 8192,
+  "hidden_size": 384,
+  "num_attention_heads": 6,
+  "intermediate_size": 1024,
+  "num_hidden_layers": 4,
+  "local_window_size": 64,
+  "max_position_embeddings": 512,
+  "torch_dtype": "float32",
+  "pad_token_id": 0,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "unk_token_id": 3
+}

generate.py

@@ -0,0 +1,62 @@
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_model
+from tokenizers import Tokenizer
+
+from model import CED
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model-dir",
+        type=Path,
+        default=Path(__file__).resolve().parent,
+    )
+    parser.add_argument("--text", default="Once upon a time, a little rabbit")
+    parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
+    parser.add_argument("--temperature", type=float, default=0.8)
+    parser.add_argument("--max-new-tokens", type=int, default=150)
+    args = parser.parse_args()
+
+    config = json.loads((args.model_dir / "config.json").read_text())
+    model = CED(
+        vocab=config["vocab_size"],
+        dim=config["hidden_size"],
+        heads=config["num_attention_heads"],
+        ff=config["intermediate_size"],
+        layers=config["num_hidden_layers"],
+        window=config["local_window_size"],
+    )
+    load_model(model, args.model_dir / "model.safetensors")
+    model = model.to(args.device).eval()
+    tokenizer = Tokenizer.from_file(str(args.model_dir / "tokenizer.json"))
+
+    ids = [config["bos_token_id"]] + tokenizer.encode(args.text).ids
+    context_length = config["max_position_embeddings"]
+    if len(ids) > context_length:
+        raise ValueError(f"Prompt exceeds the {context_length}-token context length")
+
+    torch.manual_seed(123)
+    with torch.inference_mode():
+        for _ in range(min(args.max_new_tokens, context_length - len(ids))):
+            logits = model(torch.tensor([ids], device=args.device))[0, -1].float()
+            logits[
+                [config["pad_token_id"], config["bos_token_id"], config["unk_token_id"]]
+            ] = -torch.inf
+            if args.temperature <= 0:
+                token = int(logits.argmax())
+            else:
+                token = int(torch.multinomial((logits / args.temperature).softmax(-1), 1))
+            ids.append(token)
+            if token == config["eos_token_id"]:
+                break
+
+    print(tokenizer.decode(ids))
+
+
+if __name__ == "__main__":
+    main()

model.py

@@ -0,0 +1,109 @@
+"""Tiny CED model architecture required to load model.safetensors."""
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        z = x.float()
+        return (
+            z
+            * torch.rsqrt(z.square().mean(-1, keepdim=True) + 1e-5)
+            * self.weight
+        ).to(x.dtype)
+
+
+def rope(x, pos):
+    d = x.shape[-1]
+    freq = 10000 ** (torch.arange(0, d, 2, device=x.device) / d)
+    angle = pos[:, None, :, None].float() * freq
+    a, b = x.float()[..., ::2], x.float()[..., 1::2]
+    c, s = angle.cos(), angle.sin()
+    return torch.stack((a * c - b * s, a * s + b * c), dim=-1).flatten(-2).to(x.dtype)
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads, window):
+        super().__init__()
+        self.heads = heads
+        self.window = window
+        for name in ["wq", "wkg", "wvg", "wkl", "wvl", "wo"]:
+            setattr(self, name, nn.Linear(dim, dim, bias=False))
+
+    def forward(self, u, source, pos, valid):
+        batch, length, dim = u.shape
+
+        def split(z):
+            return z.view(batch, length, self.heads, dim // self.heads).transpose(1, 2)
+
+        q = rope(split(self.wq(u)), pos)
+        kg, vg = rope(split(self.wkg(source)), pos), split(self.wvg(source))
+        kl, vl = rope(split(self.wkl(u)), pos), split(self.wvl(u))
+        distance = pos[:, :, None] - pos[:, None, :]
+        global_mask = (distance >= 0) & valid[:, :, None] & valid[:, None, :]
+        local_mask = global_mask & (distance < self.window)
+        mask = torch.cat((global_mask, local_mask), -1)[:, None]
+        k, v = torch.cat((kg, kl), 2), torch.cat((vg, vl), 2)
+        out = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=mask,
+            dropout_p=0,
+            is_causal=False,
+        )
+        return self.wo(out.transpose(1, 2).contiguous().view(batch, length, dim))
+
+
+class Block(nn.Module):
+    def __init__(self, dim, heads, ff, window):
+        super().__init__()
+        self.attn_norm, self.ffn_norm = RMSNorm(dim), RMSNorm(dim)
+        self.attn = Attention(dim, heads, window)
+        self.gate, self.up = (nn.Linear(dim, ff, bias=False) for _ in range(2))
+        self.down = nn.Linear(ff, dim, bias=False)
+
+    def forward(self, h, e, pos, valid):
+        u = self.attn_norm(h)
+        h = h + self.attn(u, u if e is None else e, pos, valid)
+        z = self.ffn_norm(h)
+        return h + self.down(F.silu(self.gate(z)) * self.up(z))
+
+
+class CED(nn.Module):
+    def __init__(self, vocab=8192, dim=384, heads=6, ff=1024, layers=4, window=64):
+        super().__init__()
+        assert dim % heads == 0 and (dim // heads) % 2 == 0
+        self.embedding = nn.Embedding(vocab, dim)
+        self.encoder = nn.ModuleList(
+            [Block(dim, heads, ff, window) for _ in range(layers)]
+        )
+        self.decoder = nn.ModuleList(
+            [Block(dim, heads, ff, window) for _ in range(layers)]
+        )
+        self.encoder_norm, self.final_norm = RMSNorm(dim), RMSNorm(dim)
+        self.head = nn.Linear(dim, vocab, bias=False)
+        self.head.weight = self.embedding.weight
+        for parameter in self.parameters():
+            if parameter.ndim > 1:
+                nn.init.normal_(parameter, std=0.02)
+
+    def forward(self, ids, positions=None, valid=None):
+        if positions is None:
+            positions = torch.arange(ids.shape[1], device=ids.device)[None].expand_as(ids)
+        if valid is None:
+            valid = torch.ones_like(ids, dtype=torch.bool)
+        h = self.embedding(ids)
+        for layer in self.encoder:
+            h = layer(h, None, positions, valid)
+        e = self.encoder_norm(h)
+        h = e
+        for layer in self.decoder:
+            h = layer(h, e, positions, valid)
+        return self.head(self.final_norm(h))

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a60e58abb8fa784857588179cba839ccc76d4d56ef297887c376f3d8a5aacf5
+size 78679560

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.0
+safetensors>=0.4
+tokenizers>=0.20