{"id":"rustyspottedcatt/rollm","name":"rollm","scope":"rustyspottedcatt","platform":"roblox","description":"A lightweight educational transformer language-model framework for Roblox Luau.","version":"0.2.1","latest":"0.2.1","versions":["0.2.0","0.2.1"],"license":"MIT","licenseRating":"safe","licenseCaveats":[],"licenseVerified":true,"dependencies":{"evaera/promise":{"version":"^4.0.0","alias":"promise"},"sleitnick/signal":{"version":"^2.0.1","alias":"signal"}},"integrity":"e4d8cbfc72d72d08d015cdf18be7db2096c4e62fe45ee5eb2b78223a18821194","likes":0,"downloads":0,"install":"forest install rustyspottedcatt/rollm","url":"https://forest.dev/p/roblox/rustyspottedcatt/rollm","files":"https://api.forest.dev/ai/package/roblox/rustyspottedcatt/rollm/files","readme":"<img width=\"1500\" height=\"500\" alt=\"image\" src=\"https://github.com/user-attachments/assets/d34446b8-0a26-4015-83cf-4ee2af32b075\" />\n\n# RoLLM\n\n[![CI](https://github.com/rustyspottedcatt/RoLLM/actions/workflows/ci.yml/badge.svg)](https://github.com/rustyspottedcatt/RoLLM/actions/workflows/ci.yml)\n[![CD](https://github.com/rustyspottedcatt/RoLLM/actions/workflows/cd.yml/badge.svg)](https://github.com/rustyspottedcatt/RoLLM/actions/workflows/cd.yml)\n[![Wally](https://img.shields.io/badge/wally-rustyspottedcatt%2Frollm-orange)](https://wally.run/package/rustyspottedcatt/rollm)\n[![License](https://img.shields.io/badge/License-MIT-blue)](LICENSE)\n[![Maintainer](https://img.shields.io/badge/maintainer-rustyspotted-blue)](https://github.com/rustyspottedcatt)\n\nA transformer language model framework for Roblox/Luau. The full stack — tokenizers, positional embeddings, multi-head attention, layer norm, feedforward blocks, Adam, cross-entropy — in pure Luau. No native code, no external runtime, no magic.\n\n---\n\n## Table of Contents\n\n- [Features](#features)\n- [Architecture](#architecture)\n- [Installation](#installation)\n- [Quick Start](#quick-start)\n- [API Reference](#api-reference)\n  - [RoLLM.new](#rollmnew)\n  - [trainModel](#trainmodel)\n  - [generate / generateTemperature](#generate--generatetemperature)\n  - [generateTopK / generateNucleus / generateWithKVCache](#generatetopk--generatenucleus--generatewithkvcache)\n  - [predict / predictTemperature](#predict--predicttemperature)\n  - [saveModel / loadModel](#savemodel--loadmodel)\n  - [getParameterCount](#getparametercount)\n- [Configuration](#configuration)\n- [TrainingOptions](#trainingoptions)\n- [Modules](#modules)\n- [Tokenizer Modes](#tokenizer-modes)\n- [Performance Notes](#performance-notes)\n- [Examples](#examples)\n- [Contributing](#contributing)\n- [Security](#security)\n- [License](#license)\n\n---\n\n## Features\n\n**Core (v0.1)**\n- **Char/BPE Tokenizers** — char builds vocab from the corpus directly; BPE loads an external JSON vocab over HTTP.\n- **Sinusoidal Positional Embeddings** — computed once at construction, not every forward pass.\n- **Multi-Head Causal Attention** — scaled dot-product with a cached causal mask.\n- **Transformer Blocks** — pre-norm LayerNorm, residual connections.\n- **Adam Optimizer** — bias-corrected first and second moment estimates.\n- **Cross-Entropy Loss** — numerically stable softmax, full backward pass.\n- **Temperature Sampling** — greedy (`temperature = 0`) or distribution-sampled decoding.\n- **Async Training** — `yieldEverySamples` stops Roblox from timing out long training runs.\n\n**Generative (v0.2)**\n- **GELU Activation** — replaces ReLU in the feedforward block; exact `tanh` formulation with a matching backward.\n- **Xavier Initialization** — all weight matrices (`Wq`, `Wk`, `Wv`, `Wo`, `W1`, `W2`) initialized with Xavier normal.\n- **Weight Tying** — set `weightTying = true` to share embedding and output projection weights, cutting parameter count.\n- **Gradient Clipping** — `gradClipNorm` in `TrainingOptions` clips the global L2 gradient norm before each Adam step.\n- **LR Scheduling** — linear warmup then cosine decay via `lrSchedule = true` and `warmupEpochs`.\n- **Top-k Sampling** — `generateTopK` samples from the top `k` logits after temperature scaling.\n- **Nucleus Sampling** — `generateNucleus` samples from the smallest set of tokens whose cumulative probability ≥ `p`.\n- **KV-Cache Generation** — `generateWithKVCache` seeds the cache with the prompt once, then appends one token at a time.\n- **AdaptiveVocab** — maps a large BPE vocabulary down to a smaller dense vocab built from actual corpus token frequencies.\n- **CorpusFetcher** — fetches and caches text corpora from HTTP URLs into DataStore so you don't re-download on every restart.\n- **Multi-Key DataStore** — `saveModel` / `loadModel` chunk the serialized weights into ≤ 3.5 MB pieces to work around DataStore limits.\n\n---\n\n## Architecture\n\n```\nInput Text\n    │\n    ▼\nTokenizer (char | bpe)  ──optional──▶  AdaptiveVocab (dense remapping)\n    │  textToTokens / tokensToText\n    ▼\nEmbedding  (vocabSize × dModel, Xavier init)  +  Sinusoidal Positions\n    │                                              ↑ shared when weightTying = true\n    ▼  ×numLayers                                  │\nTransformerBlock                                   │\n  ├─ MultiHeadAttention  (causal mask, KV-cache for inference)\n  ├─ Residual + LayerNorm\n  ├─ FeedForward  (Linear → GELU → Linear, Xavier init)\n  └─ Residual + LayerNorm\n    │\n    ▼\nFinal Projection  (dModel × vocabSize)\n    │\n    ▼\nLogits  →  greedy | temperature | top-k | nucleus sampling\n```\n\n---\n\n## Installation\n\n### Wally (recommended)\n\n```toml\n[dependencies]\nRoLLM = \"rustyspottedcatt/rollm@0.2.0\"\n```\n\n```sh\nwally install\n```\n\n### Manual (Rojo)\n\nClone the repo and point `src/` at `ReplicatedStorage/RoLLM` in your project tree, or use the provided `default.project.json`. Run `wally install` to pull in the `promise` and `signal` dependencies.\n\n```\nRoLLM/\n├── components/\n│   ├── BPETokenizer.luau\n│   ├── CharTokenizer.luau\n│   ├── Embedding.luau\n│   ├── FeedForward.luau\n│   ├── LayerNorm.luau\n│   ├── LinearAlgebra.luau\n│   ├── MultiHeadAttention.luau\n│   ├── Tokenizer.luau\n│   ├── TransformerBlock.luau\n│   └── TransformerModel.luau\n├── lib/\n│   ├── AdaptiveVocab.luau\n│   ├── CorpusFetcher.luau\n│   ├── CrossEntropyLoss.luau\n│   ├── Optimizer.luau\n│   └── types.luau\n└── init.luau\n```\n\n---\n\n## Quick Start\n\n### Basic (v0.1 style, still works)\n\n```lua\nlocal RoLLM = require(game:GetService(\"ReplicatedStorage\").RoLLM)\n\nlocal corpus = { \"hello roblox\", \"hello rollm\", \"roblox lua\" }\n\nlocal model = RoLLM.new(corpus, {\n    dModel        = 8,\n    numHeads      = 2,\n    dFF           = 16,\n    numLayers     = 1,\n    maxSeqLen     = 16,\n    tokenizerMode = \"char\",\n})\n\nmodel:trainModel(corpus, 3, 0.01)\nprint(model:generate(\"hel\", 10))\n```\n\n### Generative (v0.2)\n\n```lua\nlocal RoLLM = require(game:GetService(\"ReplicatedStorage\").RoLLM)\n\nlocal corpus = { \"press space to jump\", \"collect coins for points\", \"defeat the boss\" }\n\nlocal model = RoLLM.new(corpus, {\n    dModel        = 128,\n    numHeads      = 4,\n    dFF           = 256,\n    numLayers     = 4,\n    maxSeqLen     = 64,\n    tokenizerMode = \"char\",\n    weightTying   = true,\n})\n\nmodel:trainModel(corpus, 200, 3e-4, {\n    gradClipNorm = 1.0,\n    lrSchedule   = true,\n    warmupEpochs = 10,\n})\n\n-- Three ways to generate:\nprint(model:generateTopK(\"press\", 20, 40, 0.9))\nprint(model:generateNucleus(\"collect\", 20, 0.95, 0.85))\nprint(model:generateWithKVCache(\"defeat\", 20, 0.9, 0.7))\n\nmodel:saveModel(\"MyBotWeights\")\n```\n\n---\n\n## API Reference\n\n### `RoLLM.new`\n\n```lua\nRoLLM.new(\n    textData  : string | {string},\n    config    : TransformerConfig,\n    chunkSize : number?            -- default 500 000\n) -> LLMInstance\n```\n\nBuilds the tokenizer vocab from `textData`, constructs the model, and returns the instance. Raises if `dModel` is not divisible by `numHeads` or if required config fields are missing.\n\n---\n\n### `trainModel`\n\n```lua\nmodel:trainModel(\n    trainingData : {string},\n    epochs       : number,\n    learningRate : number,\n    options      : TrainingOptions?\n) -> {TrainingEpoch}\n```\n\nReturns a history table where each entry is `{ epoch, loss, averageLoss, samples }`.\n\n---\n\n### `generate` / `generateTemperature`\n\n```lua\nmodel:generate(inputStr: string, numTokens: number) -> string\nmodel:generateTemperature(inputStr: string, numTokens: number, temperature: number) -> string\n```\n\nGreedy decoding and temperature-sampled decoding. `temperature = 0` is greedy; higher values flatten the distribution.\n\n---\n\n### `generateTopK` / `generateNucleus` / `generateWithKVCache`\n\n```lua\nmodel:generateTopK(inputStr: string, numTokens: number, k: number, temperature: number?) -> string\nmodel:generateNucleus(inputStr: string, numTokens: number, p: number, temperature: number?) -> string\nmodel:generateWithKVCache(inputStr: string, numTokens: number, p: number?, temperature: number?) -> string\n```\n\n- **`generateTopK`** — samples from the top `k` logits after temperature scaling.\n- **`generateNucleus`** — samples from the smallest set of tokens whose cumulative probability ≥ `p`.\n- **`generateWithKVCache`** — nucleus sampling with a KV-cache. Seeds the cache with the full prompt once, then appends one token per step. Faster for long outputs because each new token only processes one row instead of the full sequence.\n\n---\n\n### `predict` / `predictTemperature`\n\n```lua\nmodel:predict(inputStr: string) -> string\nmodel:predictTemperature(inputStr: string, temperature: number) -> string\n```\n\nReturns the single next-token prediction as a string.\n\n---\n\n### `saveModel` / `loadModel`\n\n```lua\nmodel:saveModel(storeKey: string) -> boolean\nmodel:loadModel(storeKey: string) -> boolean\n```\n\nSerializes model weights to JSON, splits into ≤ 3.5 MB chunks, and stores them under `storeKey_meta` + `storeKey_chunk0`, `storeKey_chunk1`, etc. `loadModel` reads the metadata key first, then reassembles. Returns `false` on any DataStore error. Requires `DataStoreService` access.\n\n---\n\n### `getParameterCount`\n\n```lua\nmodel:getParameterCount() -> number\n```\n\nTotal trainable scalar parameters across embedding, all blocks, and the final projection.\n\n---\n\n## Configuration\n\n`TransformerConfig` fields:\n\n| Field | Type | Required | Description |\n|---|---|---|---|\n| `dModel` | `number` | ✓ | Hidden dimension. Must be divisible by `numHeads`. |\n| `numHeads` | `number` | ✓ | Number of attention heads. |\n| `dFF` | `number` | ✓ | Feedforward inner dimension. |\n| `numLayers` | `number` | ✓ | Number of transformer blocks. |\n| `maxSeqLen` | `number` | ✓ | Maximum token sequence length. Inputs longer than this are trimmed. |\n| `tokenizerMode` | `\"char\"\\|\"bpe\"` | | Default `\"char\"`. |\n| `externalVocabURL` | `string` | | BPE vocab JSON URL. Required when `tokenizerMode = \"bpe\"`. |\n| `weightTying` | `boolean` | | Share embedding and output projection weights. Default `false`. |\n| `vocabSize` | `number` | | Set automatically; do not pass manually. |\n\n---\n\n## TrainingOptions\n\n```lua\ntype TrainingOptions = {\n    batchSize                 : number?,   -- gradient accumulation batch size\n    yieldEverySamples         : number?,   -- call task.wait() every N samples\n    yieldEveryParameterUpdates: number?,   -- reserved\n    gradClipNorm              : number?,   -- max global L2 gradient norm (e.g. 1.0)\n    lrSchedule                : boolean?,  -- enable linear warmup + cosine decay\n    warmupEpochs              : number?,   -- epochs for the warmup phase\n    onEpochComplete           : ((epoch: number, total: number, avgLoss: number, elapsed: number) -> ())?,\n}\n```\n\n`yieldEverySamples = 1` is still the safest option to avoid script timeouts in long runs. Combine `gradClipNorm = 1.0` with `lrSchedule = true` for more stable training on larger models.\n\n---\n\n## Modules\n\n| Module | Path | Role |\n|---|---|---|\n| `RoLLM` | `src/init.luau` | Public API, training loop |\n| `TransformerModel` | `components/TransformerModel.luau` | Forward / backward, sampling, KV-cache |\n| `TransformerBlock` | `components/TransformerBlock.luau` | MHA + FFN + residuals + norms |\n| `MultiHeadAttention` | `components/MultiHeadAttention.luau` | Scaled dot-product, causal mask, KV-cache |\n| `FeedForward` | `components/FeedForward.luau` | Two-layer GELU FFN, Xavier init |\n| `LayerNorm` | `components/LayerNorm.luau` | Per-row layer normalization |\n| `Embedding` | `components/Embedding.luau` | Token embeddings + sinusoidal positions |\n| `LinearAlgebra` | `components/LinearAlgebra.luau` | Matrix ops (mm, softmax, xavier, GELU, …) |\n| `Tokenizer` | `components/Tokenizer.luau` | Mode factory (char / bpe) |\n| `CharTokenizer` | `components/CharTokenizer.luau` | Character-level tokenizer |\n| `BPETokenizer` | `components/BPETokenizer.luau` | BPE tokenizer with external vocab |\n| `AdaptiveVocab` | `lib/AdaptiveVocab.luau` | Remaps a large BPE vocab to a smaller dense vocab |\n| `CorpusFetcher` | `lib/CorpusFetcher.luau` | HTTP corpus fetch with DataStore sentence cache |\n| `CrossEntropyLoss` | `lib/CrossEntropyLoss.luau` | Stable softmax CE loss + backward |\n| `Optimizer` | `lib/Optimizer.luau` | Adam + grad clipping + LR scheduling |\n| `types` | `lib/types.luau` | Shared Luau type exports |\n\n---\n\n## Tokenizer Modes\n\n- **`\"char\"`** (default) — one token per character, vocab built entirely from the training text. Nothing to download, works offline.\n- **`\"bpe\"`** — Byte-Pair Encoding. Needs an external JSON vocab at `config.externalVocabURL`. Better sub-word coverage for natural language, but requires HTTP access.\n\nYou can layer **AdaptiveVocab** on top of either tokenizer to shrink a large vocab down to the tokens that actually appear in your corpus. Useful if you're loading a GPT-2–style vocab but only training on game-dialogue sentences.\n\n---\n\n## Performance Notes\n\n- **Causal mask cache** — the `seqLen × seqLen` mask is built once per distinct sequence length and reused.\n- **Position index cache** — position arrays are cached by length, never re-allocated.\n- **Training tokenization** — all samples tokenize before the epoch loop, not inside it.\n- **Row-major `mm`** — `LinearAlgebra:mm` pre-transposes the right operand so the inner loop hits sequential memory.\n- **Attention backward** — reuses the cached softmax weights rather than recomputing `math.exp`.\n- **KV-cache** — at inference time `forwardCached` grows the K/V tables one row at a time instead of rebuilding the full attention every step.\n\nRough Studio numbers to calibrate expectations:\n\n| Config | Params | Train (1 epoch / 20 samples) |\n|---|---|---|\n| `dModel=8, dFF=16, L=1` | ~500 | < 0.1s |\n| `dModel=64, dFF=128, L=2` | ~85k | ~1–2s |\n| `dModel=128, dFF=256, L=4` | ~430k | ~8–15s |\n\nSet `yieldEverySamples` to something small (1–5) for the larger configs.\n\n---\n\n## Examples\n\n| Script | Location | Description |\n|---|---|---|\n| Tiny training | `examples/TinyTraining.server.luau` | Minimal train-and-generate loop |\n| In-game chatbot | `examples/InGameChatBot.server.luau` | Chat bot responding to `!rollm <prompt>` |\n| Generative bot | `examples/GenerativeBot.server.luau` | Full v0.2 demo: weight tying, grad clipping, LR schedule, KV-cache, top-k, nucleus, DataStore save/load |\n| Smoke tests | `tests/SmokeTests.server.luau` | Roblox-runtime assertions covering v0.1 and v0.2 |\n| Unit tests | `tests/unit/` | Lune-runnable pure-math tests (run via `lune run tests/run_lune.luau` after `darklua process src dist`) |\n\n---\n\n## Contributing\n\nRead [CONTRIBUTING.md](CONTRIBUTING.md) before opening a pull request. Use the [bug report template](.github/ISSUE_TEMPLATE/bug_report.yml) for bugs and the [feature request template](.github/ISSUE_TEMPLATE/feature_request.yml) for new ideas.\n\nAll contributors are expected to follow the [Code of Conduct](CODE_OF_CONDUCT.md).\n\n---\n\n## Security\n\nReport vulnerabilities privately via [SECURITY.md](SECURITY.md).\n\n---\n\n## License\n\n[MIT License](LICENSE) — © 2025-2026 [rustyspotted](https://github.com/rustyspottedcatt)\n","readmeTruncated":false}