1 files changed, 19 insertions, 20 deletions

diff --git a/research/knobloch/cmd/train/main.go b/research/knobloch/cmd/train/main.go
index 72f7e00..ebe8c54 100644
--- a/research/knobloch/cmd/train/main.go
+++ b/research/knobloch/cmd/train/main.go
@@ -10,25 +10,24 @@ import (
 	"github.com/jonasknobloch/mbpe"
 
 	"go.jknobloc.com/x/dataset"
-	"go.jknobloc.com/x/profile"
 	"go.jknobloc.com/x/shelf"
 	"go.jknobloc.com/x/tokenizer/bpe"
 	"go.jknobloc.com/x/tokenizer/bpe/split"
 )
 
 func main() {
-	stop := profile.CPU()
+	// stop := profile.CPU()
 
-	// train()
+	train()
 	// serialize()
 
-	profile.Mem()
+	// profile.Mem()
 
-	stop()
+	// stop()
 }
 
 func train() {
-	out := shelf.Abs("results/knobloch/minipile")
+	out := shelf.Abs("results/knobloch/fineweb_verify")
 
 	if err := os.MkdirAll(out, os.ModePerm); err != nil {
 		log.Fatal(err)
@@ -44,7 +43,7 @@ func train() {
 		return m
 	}()
 
-	// mbpe.InvertWeightFunction = true
+	mbpe.InvertWeightFunction = true
 	// mbpe.UseSimpleClashes = true
 
 	newTrainer := func(segmenter mbpe.Segmenter, alpha float64) *mbpe.MBPETrainer {
@@ -58,24 +57,24 @@ func train() {
 
 		b.SetMatcher(split.NewFSA())
 
-		return mbpe.NewMBPETrainer(b, segmenter, alpha, mbpe.NewMBPE(), 1<<17, alphabet)
+		return mbpe.NewMBPETrainer(b, segmenter, alpha, mbpe.NewMBPE(), 1<<19, alphabet)
 	}
 
 	configs := []struct {
 		float64
 		string
 	}{
-		{0.0, "m000_minipile_v2"},
-		{0.1, "m010_minipile_v2"},
-		{0.2, "m020_minipile_v2"},
-		{0.3, "m030_minipile_v2"},
-		{0.4, "m040_minipile_v2"},
-		{0.5, "m050_minipile_v2"},
-		{0.6, "m060_minipile_v2"},
-		{0.7, "m070_minipile_v2"},
-		{0.8, "m080_minipile_v2"},
-		{0.9, "m090_minipile_v2"},
-		{1.0, "m100_minipile_v2"},
+		{0.0, "mi000_fineweb"},
+		{0.1, "mi010_fineweb"},
+		{0.2, "mi020_fineweb"},
+		{0.3, "mi030_fineweb"},
+		{0.4, "mi040_fineweb"},
+		{0.5, "mi050_fineweb"},
+		{0.6, "mi060_fineweb"},
+		{0.7, "mi070_fineweb"},
+		{0.8, "mi080_fineweb"},
+		{0.9, "mi090_fineweb"},
+		{1.0, "mi100_fineweb"},
 	}
 
 	for i, c := range configs {
@@ -86,7 +85,7 @@ func train() {
 		if dictErr := t.LoadDict(dict); dictErr != nil {
 			var reader dataset.Reader
 
-			if r, err := dataset.NewParquetReader(shelf.Abs("data/minipile/train")); err != nil {
+			if r, err := dataset.NewParquetReader(shelf.Abs("data/fineweb-edu/sample-10BT/train")); err != nil {
 				log.Fatal(err)
 			} else {
 				reader = r