1 files changed, 45 insertions, 0 deletions
diff --git a/tokenizer/bpe/cmd/tokenize/main.go b/tokenizer/bpe/cmd/tokenize/main.go
index aaf10e7..1b0b868 100644
--- a/tokenizer/bpe/cmd/tokenize/main.go
+++ b/tokenizer/bpe/cmd/tokenize/main.go
@@ -1,7 +1,11 @@
 package main
 
 import (
+	"flag"
 	"log"
+	"os"
+	"runtime"
+	"runtime/pprof"
 	"sync/atomic"
 	"time"
 
@@ -11,7 +15,30 @@ import (
 	"go.jknobloc.com/x/tui"
 )
 
+var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to `file`")
+var memprofile = flag.String("memprofile", "", "write memory profile to `file`")
+
 func main() {
+	flag.Parse()
+
+	if *cpuprofile != "" {
+		var file *os.File
+
+		if f, err := os.Create(*cpuprofile); err != nil {
+			log.Fatal("could not create CPU profile: ", err)
+		} else {
+			file = f
+
+			defer file.Close()
+		}
+
+		if err := pprof.StartCPUProfile(file); err != nil {
+			log.Fatal("could not start CPU profile: ", err)
+		}
+
+		defer pprof.StopCPUProfile()
+	}
+
 	reader := data()
 
 	t := tokenizer()
@@ -41,6 +68,24 @@ func main() {
 
 		n++
 	}
+
+	if *memprofile != "" {
+		var file *os.File
+
+		if f, err := os.Create(*memprofile); err != nil {
+			log.Fatal("could not create memory profile: ", err)
+		} else {
+			file = f
+
+			defer file.Close()
+		}
+
+		runtime.GC()
+
+		if err := pprof.Lookup("allocs").WriteTo(file, 0); err != nil {
+			log.Fatal("could not write memory profile: ", err)
+		}
+	}
 }
 
 func data() *dataset.ParquetReader {