Add llmc module

11 files changed, 706 insertions, 0 deletions

diff --git a/go.work b/go.work
index 801e7bf..9978332 100644
--- a/go.work
+++ b/go.work
@@ -4,6 +4,7 @@ use (
 	./dataset
 	./gpt2
 	./llm
+	./llmc
 	./mbpe
 	./onnx
 	./research/lesci
diff --git a/llmc/cmd/data/fineweb.go b/llmc/cmd/data/fineweb.go
new file mode 100644
index 0000000..cabfa74
--- /dev/null
+++ b/llmc/cmd/data/fineweb.go
@@ -0,0 +1,43 @@
+package main
+
+import (
+	"log"
+
+	"go.jknobloc.com/x/dataset"
+	"go.jknobloc.com/x/llm"
+	"go.jknobloc.com/x/llmc"
+	"go.jknobloc.com/x/tokenizer/bpe"
+)
+
+// classic + 10B: fineweb10B; sample-10BT
+// classic + 100B: fineweb100B; sample-100BT
+// edu + 10B: edu_fineweb10B; sample-10BT
+// edu + 100B: edu_fineweb100B; sample-100BT
+
+func fineWeb() {
+	var reader dataset.Reader
+
+	if r, err := dataset.NewParquetReader("artifacts/data/raw/fineweb-edu/sample-100BT/train"); err != nil {
+		log.Fatal(err)
+	} else {
+		reader = r
+	}
+
+	var tokenizer llm.Tokenizer
+
+	if t, err := bpe.NewTokenizerFromFiles("gpt2/models/base/vocab.json", "gpt2/models/base/merges.txt"); err != nil {
+		log.Fatal(err)
+	} else {
+		tokenizer = t
+	}
+
+	docs := llmc.TokenizeAll(reader, tokenizer, 50256)
+
+	if err := llmc.WriteShards("artifacts/data/llmc/edu_fineweb100B", "edu_fineweb", 100_000_000, docs); err != nil {
+		log.Fatal(err)
+	}
+
+	if err := reader.Err(); err != nil {
+		log.Fatal(err)
+	}
+}
diff --git a/llmc/cmd/data/main.go b/llmc/cmd/data/main.go
new file mode 100644
index 0000000..b12d124
--- /dev/null
+++ b/llmc/cmd/data/main.go
@@ -0,0 +1,6 @@
+package main
+
+func main() {
+	// tinyShakespeare()
+	// fineWeb()
+}
diff --git a/llmc/cmd/data/tinyshakespeare.go b/llmc/cmd/data/tinyshakespeare.go
new file mode 100644
index 0000000..b2d062c
--- /dev/null
+++ b/llmc/cmd/data/tinyshakespeare.go
@@ -0,0 +1,66 @@
+package main
+
+import (
+	"fmt"
+	"log"
+
+	"go.jknobloc.com/x/dataset"
+	"go.jknobloc.com/x/llmc"
+	"go.jknobloc.com/x/tokenizer/bpe"
+)
+
+func tinyShakespeare() {
+	var reader *dataset.FileReader
+
+	if r, err := dataset.NewFileReader("artifacts/data/raw/tinyshakespeare/input.txt", "*.txt"); err != nil {
+		log.Fatal(err)
+	} else {
+		reader = r
+	}
+
+	reader.SetDelimiters("\n\n")
+
+	var tokenizer *bpe.Tokenizer
+
+	if t, err := bpe.NewTokenizerFromFiles("gpt2/models/base/vocab.json", "gpt2/models/base/merges.txt"); err != nil {
+		log.Fatal(err)
+	} else {
+		tokenizer = t
+	}
+
+	result := make([]int64, 0)
+
+	for _, doc := range reader.Texts() {
+		ids := tokenizer.Tokenize(doc)
+
+		toks := make([]int64, len(ids)+1)
+
+		toks[0] = 50256 // end of text
+
+		for i, id := range ids {
+			toks[i+1] = int64(id)
+		}
+
+		result = append(result, toks...)
+	}
+
+	val := llmc.DataFile[int64]{
+		Model:  llmc.GPT2,
+		Tokens: result[:32768],
+	}
+
+	train := llmc.DataFile[int64]{
+		Model:  llmc.GPT2,
+		Tokens: result[32768:],
+	}
+
+	if _, err := llmc.Serialize(&val, "artifacts/data/llmc/tinyshakespeare/tiny_shakespeare_val.bin"); err != nil {
+		log.Fatal(err)
+	}
+
+	if _, err := llmc.Serialize(&train, "artifacts/data/llmc/tinyshakespeare/tiny_shakespeare_train.bin"); err != nil {
+		log.Fatal(err)
+	}
+
+	fmt.Println(val)
+}
diff --git a/llmc/cmd/debug/debug.go b/llmc/cmd/debug/debug.go
new file mode 100644
index 0000000..289e55f
--- /dev/null
+++ b/llmc/cmd/debug/debug.go
@@ -0,0 +1,49 @@
+package main
+
+import (
+	"fmt"
+	"log"
+
+	"go.jknobloc.com/x/llmc"
+	"go.jknobloc.com/x/tokenizer/bpe"
+)
+
+func main() {
+	var data llmc.DataFile[uint16]
+	var gold llmc.DataFile[uint16]
+
+	_ = must(llmc.Deserialize("artifacts/data/llmc/tinyshakespeare/tiny_shakespeare_train.bin", &data))
+	_ = must(llmc.Deserialize("artifacts/test/llmc/tinyshakespeare/tiny_shakespeare_train.bin", &gold))
+
+	t := must(bpe.NewTokenizerFromFiles("gpt2/models/base/vocab.json", "gpt2/models/base/merges.txt"))
+
+	itoa := bpe.Itoa(t)
+
+	fmt.Println(len(data.Tokens))
+	fmt.Println(len(gold.Tokens))
+
+	if len(data.Tokens) != len(gold.Tokens) {
+		log.Fatal("length mismatch")
+	}
+
+	for i, a := range data.Tokens {
+		b := gold.Tokens[i]
+
+		if a != b {
+			fmt.Printf("token mismatch: %d %d at index %d\n", a, b, i)
+
+			fmt.Println("data", itoa[int64(a)])
+			fmt.Println("gold", itoa[int64(b)])
+
+			break
+		}
+	}
+}
+
+func must[T any](v T, err error) T {
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	return v
+}
diff --git a/llmc/cmd/peek/main.go b/llmc/cmd/peek/main.go
new file mode 100644
index 0000000..caacb31
--- /dev/null
+++ b/llmc/cmd/peek/main.go
@@ -0,0 +1,52 @@
+package main
+
+import (
+	"fmt"
+	"log"
+
+	"golang.org/x/exp/constraints"
+
+	"go.jknobloc.com/x/llmc"
+	"go.jknobloc.com/x/tokenizer/bpe"
+)
+
+func main() {
+	var a llmc.DataFile[uint16]
+	var b llmc.DataFile[uint16]
+
+	_ = must(llmc.Deserialize("artifacts/data/llmc/edu_fineweb100B/edu_fineweb_val_000000.bin", &a))
+	_ = must(llmc.Deserialize("artifacts/test/llmc/edu_fineweb100B/edu_fineweb_val_000000.bin", &b))
+
+	t := must(bpe.NewTokenizerFromFiles("gpt2/models/base/vocab.json", "gpt2/models/base/merges.txt"))
+
+	s := decode(&a, 1024, t)
+	k := decode(&b, 1024, t)
+
+	fmt.Println(s)
+
+	fmt.Println()
+
+	fmt.Println(k)
+}
+
+func must[T any](v T, err error) T {
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	return v
+}
+
+func decode[T constraints.Integer](src *llmc.DataFile[T], n int, t *bpe.Tokenizer) string {
+	ids := make([]int, n)
+
+	if len(src.Tokens) < n {
+		panic("not enough tokens")
+	}
+
+	for i := range n {
+		ids[i] = int(src.Tokens[i])
+	}
+
+	return t.Decode(ids)
+}
diff --git a/llmc/datafile.go b/llmc/datafile.go
new file mode 100644
index 0000000..d68f1b5
--- /dev/null
+++ b/llmc/datafile.go
@@ -0,0 +1,152 @@
+package llmc
+
+import (
+	"encoding/binary"
+	"fmt"
+	"io"
+	"os"
+
+	"golang.org/x/exp/constraints"
+)
+
+type ModelDesc string
+
+const (
+	GPT2   ModelDesc = "gpt-2"
+	LLaMA3 ModelDesc = "llama-3"
+)
+
+type modelInfo struct {
+	magic   int32
+	version int32
+	wide    bool // true: uint32 tokens; false: uint16 tokens
+}
+
+var headerInfo = map[ModelDesc]modelInfo{
+	GPT2:   {magic: 20240520, version: 1, wide: false},
+	LLaMA3: {magic: 20240801, version: 7, wide: true},
+}
+
+type DataFile[T constraints.Integer] struct {
+	Model  ModelDesc
+	Tokens []T // token IDs
+}
+
+func Deserialize[T constraints.Integer](name string, dst *DataFile[T]) (int, error) {
+	var file *os.File
+
+	if f, err := os.Open(name); err != nil {
+		return 0, err
+	} else {
+		file = f
+	}
+
+	defer file.Close()
+
+	var header [headerWords * 4]byte
+
+	if _, err := io.ReadFull(file, header[:]); err != nil {
+		return 0, err
+	}
+
+	magic := int32(binary.LittleEndian.Uint32(header[0:]))
+	version := int32(binary.LittleEndian.Uint32(header[4:]))
+	numToks := int32(binary.LittleEndian.Uint32(header[8:]))
+
+	var model ModelDesc
+	var info modelInfo
+
+	for m, mi := range headerInfo {
+		if mi.magic == magic && mi.version == version {
+			model, info = m, mi
+			break
+		}
+	}
+
+	if model == "" {
+		return 0, fmt.Errorf("unknown magic:version %d:%d", magic, version)
+	}
+
+	tokens := make([]T, numToks)
+
+	if info.wide {
+		tokensU32 := make([]uint32, numToks)
+
+		if err := binary.Read(file, binary.LittleEndian, tokensU32); err != nil {
+			return 0, err
+		}
+
+		for i, t := range tokensU32 {
+			tokens[i] = T(t)
+		}
+	} else {
+		tokensU16 := make([]uint16, numToks)
+
+		if err := binary.Read(file, binary.LittleEndian, tokensU16); err != nil {
+			return 0, err
+		}
+
+		for i, t := range tokensU16 {
+			tokens[i] = T(t)
+		}
+	}
+
+	dst.Model = model
+	dst.Tokens = tokens
+
+	return int(numToks), nil
+}
+
+func Serialize[T constraints.Integer](src *DataFile[T], name string) (int, error) {
+	info, ok := headerInfo[src.Model]
+
+	if !ok {
+		return 0, fmt.Errorf("unknown model descriptor %q", src.Model)
+	}
+
+	h := Header{
+		Magic:   info.magic,
+		Version: info.version,
+		NumToks: int32(len(src.Tokens)),
+	}
+
+	var file *os.File
+
+	if f, err := os.Create(name); err != nil {
+		return 0, err
+	} else {
+		file = f
+	}
+
+	defer file.Close()
+
+	header := h.encode()
+
+	if _, err := file.Write(header[:]); err != nil {
+		return 0, err
+	}
+
+	if info.wide {
+		tokensU32 := make([]uint32, len(src.Tokens))
+
+		for i, t := range src.Tokens {
+			tokensU32[i] = uint32(t)
+		}
+
+		if err := binary.Write(file, binary.LittleEndian, tokensU32); err != nil {
+			return 0, err
+		}
+	} else {
+		tokensU16 := make([]uint16, len(src.Tokens))
+
+		for i, t := range src.Tokens {
+			tokensU16[i] = uint16(t)
+		}
+
+		if err := binary.Write(file, binary.LittleEndian, tokensU16); err != nil {
+			return 0, err
+		}
+	}
+
+	return len(src.Tokens), nil
+}
diff --git a/llmc/go.mod b/llmc/go.mod
new file mode 100644
index 0000000..cc08275
--- /dev/null
+++ b/llmc/go.mod
@@ -0,0 +1,50 @@
+module go.jknobloc.com/x/llmc
+
+go 1.25.0
+
+require (
+	go.jknobloc.com/x/dataset v0.0.0-20260417140714-3fe37224ccc2
+	go.jknobloc.com/x/llm v0.0.0-20260417140714-3fe37224ccc2
+	go.jknobloc.com/x/tokenizer v0.0.0-20260417140714-3fe37224ccc2
+)
+
+require (
+	codeberg.org/go-fonts/liberation v0.5.0 // indirect
+	codeberg.org/go-latex/latex v0.2.0 // indirect
+	codeberg.org/go-pdf/fpdf v0.11.1 // indirect
+	git.sr.ht/~sbinet/gg v0.7.0 // indirect
+	github.com/ajstarks/svgo v0.0.0-20211024235047-1546f124cd8b // indirect
+	github.com/andybalholm/brotli v1.2.0 // indirect
+	github.com/apache/arrow-go/v18 v18.5.1 // indirect
+	github.com/apache/thrift v0.22.0 // indirect
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/goccy/go-json v0.10.5 // indirect
+	github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0 // indirect
+	github.com/golang/snappy v1.0.0 // indirect
+	github.com/google/flatbuffers v25.12.19+incompatible // indirect
+	github.com/google/uuid v1.6.0 // indirect
+	github.com/jonasknobloch/mbpe v0.1.1 // indirect
+	github.com/klauspost/asmfmt v1.3.2 // indirect
+	github.com/klauspost/compress v1.18.3 // indirect
+	github.com/klauspost/cpuid/v2 v2.3.0 // indirect
+	github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 // indirect
+	github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 // indirect
+	github.com/pierrec/lz4/v4 v4.1.25 // indirect
+	github.com/zeebo/xxh3 v1.1.0 // indirect
+	go.jknobloc.com/x/tui v0.0.0-20260324194423-87bbece7e040 // indirect
+	golang.org/x/exp v0.0.0-20260112195511-716be5621a96 // indirect
+	golang.org/x/image v0.37.0 // indirect
+	golang.org/x/mod v0.33.0 // indirect
+	golang.org/x/net v0.50.0 // indirect
+	golang.org/x/sync v0.20.0 // indirect
+	golang.org/x/sys v0.41.0 // indirect
+	golang.org/x/telemetry v0.0.0-20260209163413-e7419c687ee4 // indirect
+	golang.org/x/text v0.35.0 // indirect
+	golang.org/x/tools v0.42.0 // indirect
+	golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect
+	gonum.org/v1/gonum v0.17.0 // indirect
+	gonum.org/v1/plot v0.16.0 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20251029180050-ab9386a59fda // indirect
+	google.golang.org/grpc v1.78.0 // indirect
+	google.golang.org/protobuf v1.36.11 // indirect
+)
diff --git a/llmc/go.sum b/llmc/go.sum
new file mode 100644
index 0000000..93ca5c2
--- /dev/null
+++ b/llmc/go.sum
@@ -0,0 +1,153 @@
+codeberg.org/go-fonts/dejavu v0.4.0 h1:2yn58Vkh4CFK3ipacWUAIE3XVBGNa0y1bc95Bmfx91I=
+codeberg.org/go-fonts/dejavu v0.4.0/go.mod h1:abni088lmhQJvso2Lsb7azCKzwkfcnttl6tL1UTWKzg=
+codeberg.org/go-fonts/latin-modern v0.4.0 h1:vkRCc1y3whKA7iL9Ep0fSGVuJfqjix0ica9UflHORO8=
+codeberg.org/go-fonts/latin-modern v0.4.0/go.mod h1:BF68mZznJ9QHn+hic9ks2DaFl4sR5YhfM6xTYaP9vNw=
+codeberg.org/go-fonts/liberation v0.5.0 h1:SsKoMO1v1OZmzkG2DY+7ZkCL9U+rrWI09niOLfQ5Bo0=
+codeberg.org/go-fonts/liberation v0.5.0/go.mod h1:zS/2e1354/mJ4pGzIIaEtm/59VFCFnYC7YV6YdGl5GU=
+codeberg.org/go-latex/latex v0.2.0 h1:Ol/a6VHY06N+5gPfewswymoRb5ZcKDXWVaVegcx4hbI=
+codeberg.org/go-latex/latex v0.2.0/go.mod h1:VJAwQir7/T8LZxj7xAPivISKiVOwkMpQ8bTuPQ31X0Y=
+codeberg.org/go-pdf/fpdf v0.11.1 h1:U8+coOTDVLxHIXZgGvkfQEi/q0hYHYvEHFuGNX2GzGs=
+codeberg.org/go-pdf/fpdf v0.11.1/go.mod h1:Y0DGRAdZ0OmnZPvjbMp/1bYxmIPxm0ws4tfoPOc4LjU=
+git.sr.ht/~sbinet/cmpimg v0.1.0 h1:E0zPRk2muWuCqSKSVZIWsgtU9pjsw3eKHi8VmQeScxo=
+git.sr.ht/~sbinet/cmpimg v0.1.0/go.mod h1:FU12psLbF4TfNXkKH2ZZQ29crIqoiqTZmeQ7dkp/pxE=
+git.sr.ht/~sbinet/gg v0.7.0 h1:YmNf7YKd7diDMTPm86hZa1EM3pbkOyD/zzjl0LZUdNM=
+git.sr.ht/~sbinet/gg v0.7.0/go.mod h1:VYeli15tpMM4EvqlivlVbbyvWZlOU+EZn4XZmfBGUdM=
+github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
+github.com/ajstarks/deck v0.0.0-20200831202436-30c9fc6549a9/go.mod h1:JynElWSGnm/4RlzPXRlREEwqTHAN3T56Bv2ITsFT3gY=
+github.com/ajstarks/deck/generate v0.0.0-20210309230005-c3f852c02e19/go.mod h1:T13YZdzov6OU0A1+RfKZiZN9ca6VeKdBdyDV+BY97Tk=
+github.com/ajstarks/svgo v0.0.0-20211024235047-1546f124cd8b h1:slYM766cy2nI3BwyRiyQj/Ud48djTMtMebDqepE95rw=
+github.com/ajstarks/svgo v0.0.0-20211024235047-1546f124cd8b/go.mod h1:1KcenG0jGWcpt8ov532z81sp/kMMUG485J2InIOyADM=
+github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwToPjQ=
+github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY=
+github.com/apache/arrow-go/v18 v18.5.1 h1:yaQ6zxMGgf9YCYw4/oaeOU3AULySDlAYDOcnr4LdHdI=
+github.com/apache/arrow-go/v18 v18.5.1/go.mod h1:OCCJsmdq8AsRm8FkBSSmYTwL/s4zHW9CqxeBxEytkNE=
+github.com/apache/thrift v0.22.0 h1:r7mTJdj51TMDe6RtcmNdQxgn9XcyfGDOzegMDRg47uc=
+github.com/apache/thrift v0.22.0/go.mod h1:1e7J/O1Ae6ZQMTYdy9xa3w9k+XHWPfRvdPyJeynQ+/g=
+github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
+github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
+github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dlclark/regexp2 v1.11.5 h1:Q/sSnsKerHeCkc/jSTNq1oCm7KiVgUMZRDUoRu0JQZQ=
+github.com/dlclark/regexp2 v1.11.5/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
+github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
+github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
+github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
+github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
+github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4=
+github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M=
+github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0 h1:DACJavvAHhabrF08vX0COfcOBJRhZ8lUbR+ZWIs0Y5g=
+github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0/go.mod h1:E/TSTwGwJL78qG/PmXZO1EjYhfJinVAhrmmHX6Z8B9k=
+github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
+github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
+github.com/golang/snappy v1.0.0 h1:Oy607GVXHs7RtbggtPBnr2RmDArIsAefDwvrdWvRhGs=
+github.com/golang/snappy v1.0.0/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
+github.com/google/flatbuffers v25.12.19+incompatible h1:haMV2JRRJCe1998HeW/p0X9UaMTK6SDo0ffLn2+DbLs=
+github.com/google/flatbuffers v25.12.19+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8=
+github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
+github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/jonasknobloch/mbpe v0.1.1 h1:eXUrMdM7Wt6kPTA8cyX+t1PsInk3QIXpt430pE8CJEs=
+github.com/jonasknobloch/mbpe v0.1.1/go.mod h1:2qW/5BfAu7GKAXXNYJzMrIPYk/8z5+/TbabaAb+vPv8=
+github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
+github.com/klauspost/asmfmt v1.3.2 h1:4Ri7ox3EwapiOjCki+hw14RyKk201CN4rzyCJRFLpK4=
+github.com/klauspost/asmfmt v1.3.2/go.mod h1:AG8TuvYojzulgDAMCnYn50l/5QV3Bs/tp6j0HLHbNSE=
+github.com/klauspost/compress v1.18.3 h1:9PJRvfbmTabkOX8moIpXPbMMbYN60bWImDDU7L+/6zw=
+github.com/klauspost/compress v1.18.3/go.mod h1:R0h/fSBs8DE4ENlcrlib3PsXS61voFxhIs2DeRhCvJ4=
+github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y=
+github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
+github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs=
+github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY=
+github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 h1:+n/aFZefKZp7spd8DFdX7uMikMLXX4oubIzJF4kv/wI=
+github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3/go.mod h1:RagcQ7I8IeTMnF8JTXieKnO4Z6JCsikNEzj0DwauVzE=
+github.com/pierrec/lz4/v4 v4.1.25 h1:kocOqRffaIbU5djlIBr7Wh+cx82C0vtFb0fOurZHqD0=
+github.com/pierrec/lz4/v4 v4.1.25/go.mod h1:EoQMVJgeeEOMsCqCzqFm2O0cJvljX2nGZjcRIPL34O4=
+github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
+github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
+github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
+github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
+github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
+github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU=
+github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
+github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
+github.com/zeebo/assert v1.3.0 h1:g7C04CbJuIDKNPFHmsk4hwZDO5O+kntRxzaUoNXj+IQ=
+github.com/zeebo/assert v1.3.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0=
+github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs=
+github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s=
+go.jknobloc.com/x/dataset v0.0.0-20260417140714-3fe37224ccc2 h1:f7vUIPq1vftlvOtaoJ0Kw3jXtYoTfeQG4wxOkJaVJtM=
+go.jknobloc.com/x/dataset v0.0.0-20260417140714-3fe37224ccc2/go.mod h1:UZypBoGqi23LRt29l2z9NtVD2pKBHd4unV/sztHRz4o=
+go.jknobloc.com/x/llm v0.0.0-20260417140714-3fe37224ccc2 h1:AkrFr6n/TytOh2TIxO8fNNmT7BhG93V59OvcFDEDWjI=
+go.jknobloc.com/x/llm v0.0.0-20260417140714-3fe37224ccc2/go.mod h1:EaCZkdjylT99IfbmDpKeoAmDtaAFLodiNX1mUZpjuRM=
+go.jknobloc.com/x/tokenizer v0.0.0-20260417140714-3fe37224ccc2 h1:rmyHxhc8zMcjWSYT8axgR8tqmItCpybEkq4M08bfVNE=
+go.jknobloc.com/x/tokenizer v0.0.0-20260417140714-3fe37224ccc2/go.mod h1:A6usfKV6fYSVHGvEMO7xMMIshky4NRCBzLhs36W7KC8=
+go.jknobloc.com/x/tui v0.0.0-20260324194423-87bbece7e040 h1:7Ago/qaKuXIJ7/+/XDRO8Yb/4VWzBLKEcbAq2jD5k3Q=
+go.jknobloc.com/x/tui v0.0.0-20260324194423-87bbece7e040/go.mod h1:kblmBsWO7WlRCo4sWpstaVV9ukNfYP3PeAPFhcoJvKA=
+go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
+go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
+go.opentelemetry.io/otel v1.38.0 h1:RkfdswUDRimDg0m2Az18RKOsnI8UDzppJAtj01/Ymk8=
+go.opentelemetry.io/otel v1.38.0/go.mod h1:zcmtmQ1+YmQM9wrNsTGV/q/uyusom3P8RxwExxkZhjM=
+go.opentelemetry.io/otel/metric v1.38.0 h1:Kl6lzIYGAh5M159u9NgiRkmoMKjvbsKtYRwgfrA6WpA=
+go.opentelemetry.io/otel/metric v1.38.0/go.mod h1:kB5n/QoRM8YwmUahxvI3bO34eVtQf2i4utNVLr9gEmI=
+go.opentelemetry.io/otel/sdk v1.38.0 h1:l48sr5YbNf2hpCUj/FoGhW9yDkl+Ma+LrVl8qaM5b+E=
+go.opentelemetry.io/otel/sdk v1.38.0/go.mod h1:ghmNdGlVemJI3+ZB5iDEuk4bWA3GkTpW+DOoZMYBVVg=
+go.opentelemetry.io/otel/sdk/metric v1.38.0 h1:aSH66iL0aZqo//xXzQLYozmWrXxyFkBJ6qT5wthqPoM=
+go.opentelemetry.io/otel/sdk/metric v1.38.0/go.mod h1:dg9PBnW9XdQ1Hd6ZnRz689CbtrUp0wMMs9iPcgT9EZA=
+go.opentelemetry.io/otel/trace v1.38.0 h1:Fxk5bKrDZJUH+AMyyIXGcFAPah0oRcT+LuNtJrmcNLE=
+go.opentelemetry.io/otel/trace v1.38.0/go.mod h1:j1P9ivuFsTceSWe1oY+EeW3sc+Pp42sO++GHkg4wwhs=
+golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
+golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
+golang.org/x/exp v0.0.0-20260112195511-716be5621a96 h1:Z/6YuSHTLOHfNFdb8zVZomZr7cqNgTJvA8+Qz75D8gU=
+golang.org/x/exp v0.0.0-20260112195511-716be5621a96/go.mod h1:nzimsREAkjBCIEFtHiYkrJyT+2uy9YZJB7H1k68CXZU=
+golang.org/x/image v0.37.0 h1:ZiRjArKI8GwxZOoEtUfhrBtaCN+4b/7709dlT6SSnQA=
+golang.org/x/image v0.37.0/go.mod h1:/3f6vaXC+6CEanU4KJxbcUZyEePbyKbaLoDOe4ehFYY=
+golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
+golang.org/x/mod v0.33.0 h1:tHFzIWbBifEmbwtGz65eaWyGiGZatSrT9prnU8DbVL8=
+golang.org/x/mod v0.33.0/go.mod h1:swjeQEj+6r7fODbD2cqrnje9PnziFuw4bmLbBZFrQ5w=
+golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
+golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
+golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
+golang.org/x/net v0.50.0 h1:ucWh9eiCGyDR3vtzso0WMQinm2Dnt8cFMuQa9K33J60=
+golang.org/x/net v0.50.0/go.mod h1:UgoSli3F/pBgdJBHCTc+tp3gmrU4XswgGRgtnwWTfyM=
+golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
+golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
+golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k=
+golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
+golang.org/x/telemetry v0.0.0-20260209163413-e7419c687ee4 h1:bTLqdHv7xrGlFbvf5/TXNxy/iUwwdkjhqQTJDjW7aj0=
+golang.org/x/telemetry v0.0.0-20260209163413-e7419c687ee4/go.mod h1:g5NllXBEermZrmR51cJDQxmJUHUOfRAaNyWBM+R+548=
+golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/text v0.35.0 h1:JOVx6vVDFokkpaq1AEptVzLTpDe9KGpj5tR4/X+ybL8=
+golang.org/x/text v0.35.0/go.mod h1:khi/HExzZJ2pGnjenulevKNX1W67CUy0AsXcNubPGCA=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
+golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
+golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0=
+golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k=
+golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0=
+golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da h1:noIWHXmPHxILtqtCOPIhSt0ABwskkZKjD3bXGnZGpNY=
+golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da/go.mod h1:NDW/Ps6MPRej6fsCIbMTohpP40sJ/P/vI1MoTEGwX90=
+gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4=
+gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E=
+gonum.org/v1/plot v0.16.0 h1:dK28Qx/Ky4VmPUN/2zeW0ELyM6ucDnBAj5yun7M9n1g=
+gonum.org/v1/plot v0.16.0/go.mod h1:Xz6U1yDMi6Ni6aaXILqmVIb6Vro8E+K7Q/GeeH+Pn0c=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20251029180050-ab9386a59fda h1:i/Q+bfisr7gq6feoJnS/DlpdwEL4ihp41fvRiM3Ork0=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20251029180050-ab9386a59fda/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk=
+google.golang.org/grpc v1.78.0 h1:K1XZG/yGDJnzMdd/uZHAkVqJE+xIDOcmdSFZkBUicNc=
+google.golang.org/grpc v1.78.0/go.mod h1:I47qjTo4OKbMkjA/aOOwxDIiPSBofUtQUI5EfpWvW7U=
+google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE=
+google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+honnef.co/go/tools v0.1.3/go.mod h1:NgwopIslSNH47DimFoV78dnkksY2EFtX0ajyb3K/las=
+rsc.io/pdf v0.1.1 h1:k1MczvYDUvJBe93bYd7wrZLLUEcLZAuF824/I4e5Xr4=
+rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4=
diff --git a/llmc/header.go b/llmc/header.go
new file mode 100644
index 0000000..a154ae5
--- /dev/null
+++ b/llmc/header.go
@@ -0,0 +1,21 @@
+package llmc
+
+import "encoding/binary"
+
+const headerWords = 256
+
+type Header struct {
+	Magic   int32
+	Version int32
+	NumToks int32
+}
+
+func (h Header) encode() [headerWords * 4]byte {
+	var buffer [headerWords * 4]byte
+
+	binary.LittleEndian.PutUint32(buffer[0:], uint32(h.Magic))
+	binary.LittleEndian.PutUint32(buffer[4:], uint32(h.Version))
+	binary.LittleEndian.PutUint32(buffer[8:], uint32(h.NumToks))
+
+	return buffer
+}
diff --git a/llmc/tokenize.go b/llmc/tokenize.go
new file mode 100644
index 0000000..8237af9
--- /dev/null
+++ b/llmc/tokenize.go
@@ -0,0 +1,113 @@
+package llmc
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"runtime"
+	"sync"
+
+	"go.jknobloc.com/x/dataset"
+	"go.jknobloc.com/x/llm"
+)
+
+func TokenizeAll(reader dataset.Reader, tok llm.Tokenizer, eot int) <-chan []uint32 {
+	out := make(chan []uint32, 256)
+
+	go func() {
+		defer close(out)
+
+		sem := make(chan struct{}, max(1, runtime.NumCPU()-1))
+
+		var wg sync.WaitGroup
+
+		for _, text := range reader.Texts() {
+			sem <- struct{}{}
+
+			wg.Add(1)
+
+			go func(t string) {
+				defer func() { <-sem; wg.Done() }()
+
+				tokens := tok.Tokenize(t)
+
+				tokensU32 := make([]uint32, len(tokens)+1)
+
+				tokensU32[0] = uint32(eot)
+
+				for i, id := range tokens {
+					tokensU32[i+1] = uint32(id)
+				}
+
+				out <- tokensU32
+			}(text)
+		}
+
+		wg.Wait()
+	}()
+
+	return out
+}
+
+func WriteShards(name, data string, shardSize int, docs <-chan []uint32) error {
+	if err := os.MkdirAll(name, os.ModePerm); err != nil {
+		return err
+	}
+
+	buffer := make([]uint32, 0, shardSize)
+
+	shardIdx := 0
+
+	flush := func() error {
+		split := "train"
+
+		if shardIdx == 0 {
+			split = "val"
+		}
+
+		path := filepath.Join(name, fmt.Sprintf("%s_%s_%06d.bin", data, split, shardIdx))
+
+		d := DataFile[uint32]{
+			Model:  GPT2,
+			Tokens: buffer,
+		}
+
+		if n, err := Serialize(&d, path); err != nil {
+			return err
+		} else {
+			fmt.Printf("wrote %s (%d tokens)\n", filepath.Base(path), n)
+		}
+
+		buffer = buffer[:0]
+
+		shardIdx++
+
+		return nil
+	}
+
+	for tokens := range docs {
+		for len(tokens) > 0 {
+			space := shardSize - len(buffer)
+
+			if space >= len(tokens) {
+				buffer = append(buffer, tokens...)
+
+				break
+			}
+
+			buffer = append(buffer, tokens[:space]...)
+
+			tokens = tokens[space:] // carry remainder
+
+			if err := flush(); err != nil {
+				return err
+			}
+		}
+	}
+
+	if len(buffer) > 0 {
+		return flush()
+	}
+
+	return nil
+}