From babd70eb9d4e4b028c0eb8c32e6bed65dafa0481 Mon Sep 17 00:00:00 2001 From: Jonas Knobloch Date: Wed, 22 Apr 2026 15:36:58 +0200 Subject: Add llmc module --- llmc/cmd/data/fineweb.go | 43 +++++++++++ llmc/cmd/data/main.go | 6 ++ llmc/cmd/data/tinyshakespeare.go | 66 +++++++++++++++++ llmc/cmd/debug/debug.go | 49 +++++++++++++ llmc/cmd/peek/main.go | 52 +++++++++++++ llmc/datafile.go | 152 ++++++++++++++++++++++++++++++++++++++ llmc/go.mod | 50 +++++++++++++ llmc/go.sum | 153 +++++++++++++++++++++++++++++++++++++++ llmc/header.go | 21 ++++++ llmc/tokenize.go | 113 +++++++++++++++++++++++++++++ 10 files changed, 705 insertions(+) create mode 100644 llmc/cmd/data/fineweb.go create mode 100644 llmc/cmd/data/main.go create mode 100644 llmc/cmd/data/tinyshakespeare.go create mode 100644 llmc/cmd/debug/debug.go create mode 100644 llmc/cmd/peek/main.go create mode 100644 llmc/datafile.go create mode 100644 llmc/go.mod create mode 100644 llmc/go.sum create mode 100644 llmc/header.go create mode 100644 llmc/tokenize.go (limited to 'llmc') 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 +} -- cgit v1.2.3