Limit splitting to single delimiter

2 files changed, 11 insertions, 26 deletions

diff --git a/dataset/file.go b/dataset/file.go
index 839d79c..1d6df56 100644
--- a/dataset/file.go
+++ b/dataset/file.go
@@ -11,9 +11,9 @@ import (
 )
 
 type FileReader struct {
-	shards     []string
-	delimiters []string // TODO single delimiter
-	err        error
+	shards    []string
+	delimiter string
+	err       error
 }
 
 func NewFileReader(name, pattern string) (*FileReader, error) {
@@ -32,13 +32,13 @@ func NewFileReader(name, pattern string) (*FileReader, error) {
 	slices.Sort(shards)
 
 	return &FileReader{
-		shards:     shards,
-		delimiters: []string{"\r\n", "\n"},
+		shards:    shards,
+		delimiter: "\n",
 	}, nil
 }
 
-func (f *FileReader) SetDelimiters(delimiters ...string) *FileReader {
-	f.delimiters = delimiters
+func (f *FileReader) SetDelimiter(delimiter string) *FileReader {
+	f.delimiter = delimiter
 
 	return f
 }
@@ -81,15 +81,8 @@ func (f *FileReader) read(name string) iter.Seq[string] {
 
 		defer file.Close()
 
-		delimiters := make([][]byte, len(f.delimiters))
-
-		for i, d := range f.delimiters {
-			delimiters[i] = []byte(d)
-		}
-
-		slices.SortFunc(delimiters, func(a, b []byte) int {
-			return len(b) - len(a)
-		})
+		d := []byte(f.delimiter)
+		l := len(d)
 
 		scanner := bufio.NewScanner(file)
 
@@ -102,15 +95,7 @@ func (f *FileReader) read(name string) iter.Seq[string] {
 				return 0, nil, nil
 			}
 
-			i, l := -1, -1
-
-			for _, d := range delimiters {
-				if j := bytes.Index(data, d); j >= 0 && (i < 0 || j < i) {
-					i, l = j, len(d)
-				}
-			}
-
-			if i >= 0 {
+			if i := bytes.Index(data, d); i >= 0 {
 				return i + l, data[:i+l], nil // TODO option to consume delimiter
 			}
 
diff --git a/llmc/cmd/data/tinyshakespeare.go b/llmc/cmd/data/tinyshakespeare.go
index d71151b..960dca6 100644
--- a/llmc/cmd/data/tinyshakespeare.go
+++ b/llmc/cmd/data/tinyshakespeare.go
@@ -19,7 +19,7 @@ func tinyShakespeare() {
 		reader = r
 	}
 
-	reader.SetDelimiters("\n\n")
+	reader.SetDelimiter("\n\n")
 
 	var tokenizer *bpe.Tokenizer