diff options
| -rw-r--r-- | llmc/cmd/data/fineweb.go | 2 | ||||
| -rw-r--r-- | llmc/tokenize.go | 24 |
2 files changed, 14 insertions, 12 deletions
diff --git a/llmc/cmd/data/fineweb.go b/llmc/cmd/data/fineweb.go index a1bf7df..599089d 100644 --- a/llmc/cmd/data/fineweb.go +++ b/llmc/cmd/data/fineweb.go @@ -37,7 +37,7 @@ func fineWeb() { docs := llmc.TokenizeAll(reader, tokenizer, 50256) - if err := llmc.WriteShards(shelf.Abs("llmc/edu_fineweb100B"), "edu_fineweb", 100_000_000, docs); err != nil { + if _, err := llmc.WriteShards(shelf.Abs("llmc/edu_fineweb100B"), "edu_fineweb", 100_000_000, docs); err != nil { log.Fatal(err) } diff --git a/llmc/tokenize.go b/llmc/tokenize.go index 1863157..e90de58 100644 --- a/llmc/tokenize.go +++ b/llmc/tokenize.go @@ -34,38 +34,38 @@ func TokenizeAll(reader dataset.Reader, tok llm.Tokenizer, eot int) <-chan []uin }) } -func WriteShards(path, name string, size int, docs <-chan []uint32) error { +func WriteShards(path, name string, size int, docs <-chan []uint32) (int, error) { if err := os.MkdirAll(path, os.ModePerm); err != nil { - return err + return 0, err } buffer := make([]uint32, 0, size) - shardIdx := 0 + n := 0 flush := func() error { split := "train" - if shardIdx == 0 { + if n == 0 { split = "val" } - shardName := filepath.Join(path, fmt.Sprintf("%s_%s_%06d.bin", name, split, shardIdx)) + shardName := filepath.Join(path, fmt.Sprintf("%s_%s_%06d.bin", name, split, n)) d := DataFile[uint32]{ Model: GPT2, Tokens: buffer, } - if n, err := Serialize(&d, shardName); err != nil { + if numTokens, err := Serialize(&d, shardName); err != nil { return err } else { - fmt.Printf("wrote %s (%d tokens)\n", filepath.Base(shardName), n) + fmt.Printf("wrote %s (%d tokens)\n", filepath.Base(shardName), numTokens) } buffer = buffer[:0] - shardIdx++ + n++ return nil } @@ -85,14 +85,16 @@ func WriteShards(path, name string, size int, docs <-chan []uint32) error { tokens = tokens[space:] // carry remainder if err := flush(); err != nil { - return err + return n, err } } } if len(buffer) > 0 { - return flush() + if err := flush(); err != nil { + return n, err + } } - return nil + return n, nil } |