WIPwip-frequency

1 files changed, 220 insertions, 0 deletions

diff --git a/research/knobloch/cmd/tables/main.go b/research/knobloch/cmd/tables/main.go
new file mode 100644
index 0000000..27c3f0c
--- /dev/null
+++ b/research/knobloch/cmd/tables/main.go
@@ -0,0 +1,220 @@
+// Command tables writes the tokenizer characterisation CSVs described in
+// FREQUENCIES.md.
+//
+// Coverage is the full grid: every alignment level in both directions, at every
+// vocabulary size we have tokenizers for. Statistics are over the MiniPile train
+// dictionary.
+//
+// Encoding the dictionary dominates the runtime and the encodes are independent,
+// so they run on a worker pool. One vocabulary size is processed at a time,
+// since table C compares tokenizers within a size and the comparison inputs for
+// a whole size have to be resident together.
+package main
+
+import (
+	"flag"
+	"fmt"
+	"log"
+	"runtime"
+	"sync"
+
+	"github.com/jonasknobloch/mbpe"
+	"go.jknobloc.com/x/research/knobloch"
+	"go.jknobloc.com/x/shelf"
+)
+
+var (
+	vocabSizes = []int{8192, 16384, 32768, 50256, 100512}
+	alphas     = []int{0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100}
+)
+
+// name is the tokenizer directory suffix, e.g. m050 or mi050.
+func name(alpha int, inverted bool) string {
+	if inverted {
+		return fmt.Sprintf("mi%03d", alpha)
+	}
+
+	return fmt.Sprintf("m%03d", alpha)
+}
+
+func spec(size, alpha int, inverted bool) knobloch.TableSpec {
+	n := name(alpha, inverted)
+
+	return knobloch.TableSpec{
+		Name:      n,
+		Alignment: fmt.Sprintf("%.1f", float64(alpha)/100),
+		Inverted:  inverted,
+		VocabSize: size,
+		Dir:       shelf.Item(fmt.Sprintf("tokenizers/minipile/tokenizer_gpt2_%d_%s_minipile", size, n)),
+	}
+}
+
+func main() {
+	dict := flag.String("dict", "results/knobloch/minipile/dict.txt", "shelf-relative dictionary")
+	prefix := flag.String("prefix", "table", "output file prefix")
+	workers := flag.Int("workers", runtime.NumCPU(), "parallel encodes")
+	only := flag.Int("only", 0, "restrict to a single vocabulary size, for smoke runs")
+
+	flag.Parse()
+
+	sizes := vocabSizes
+
+	if *only != 0 {
+		sizes = []int{*only}
+	}
+
+	d := mbpe.NewDict()
+
+	if err := d.Load(shelf.Abs(shelf.Item(*dict))); err != nil {
+		log.Fatal(err)
+	}
+
+	items := d.Items()
+
+	log.Printf("dictionary: %d pre-token types", len(items))
+
+	var (
+		rowsA []knobloch.TableARow
+		rowsB []knobloch.TableBRow
+
+		aligned    []knobloch.TableCRow
+		inverted   []knobloch.TableCRow
+		invVsAlign []knobloch.TableCRow
+	)
+
+	for _, size := range sizes {
+		// every tokenizer at this size, aligned then inverted
+		var specs []knobloch.TableSpec
+
+		for _, inv := range []bool{false, true} {
+			for _, a := range alphas {
+				specs = append(specs, spec(size, a, inv))
+			}
+		}
+
+		dirs := make([]shelf.Item, 0, len(specs))
+
+		for _, s := range specs {
+			dirs = append(dirs, s.Dir)
+		}
+
+		// table B measures sharing against the intersection of every vocabulary
+		// at this size, so the baseline is just another tokenizer
+		shared, err := knobloch.VocabIntersection(dirs)
+
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		log.Printf("vocab %d: %d tokenizers, intersection %d tokens", size, len(specs), len(shared))
+
+		encoded := make([]*knobloch.Encoded, len(specs))
+
+		var wg sync.WaitGroup
+
+		queue := make(chan int)
+
+		for i := 0; i < *workers; i++ {
+			wg.Add(1)
+
+			go func() {
+				defer wg.Done()
+
+				for idx := range queue {
+					e, err := knobloch.EncodeDict(specs[idx], items)
+
+					if err != nil {
+						log.Fatal(err)
+					}
+
+					encoded[idx] = e
+				}
+			}()
+		}
+
+		for i := range specs {
+			queue <- i
+		}
+
+		close(queue)
+		wg.Wait()
+
+		log.Printf("vocab %d: encoded", size)
+
+		// index by (alpha, direction) for the pairings below
+		at := func(alpha int, inv bool) *knobloch.Encoded {
+			for i, s := range specs {
+				if s.Inverted == inv && s.Name == name(alpha, inv) {
+					return encoded[i]
+				}
+			}
+
+			log.Fatalf("missing encode for %d/%s", size, name(alpha, inv))
+
+			return nil
+		}
+
+		base := at(0, false)
+
+		for _, e := range encoded {
+			rowsA = append(rowsA, knobloch.BuildTableA(e))
+
+			b, err := knobloch.BuildTableB(e, shared)
+
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			rowsB = append(rowsB, b)
+		}
+
+		for _, a := range alphas {
+			al := at(a, false)
+			inv := at(a, true)
+
+			// aligned vs the alpha=0 baseline
+			r, err := knobloch.BuildTableC(al, base, items)
+
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			aligned = append(aligned, r)
+
+			// inverted vs the alpha=0 baseline
+			r, err = knobloch.BuildTableC(inv, base, items)
+
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			inverted = append(inverted, r)
+
+			// inverted vs aligned at matching alpha
+			r, err = knobloch.BuildTableC(inv, al, items)
+
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			invVsAlign = append(invVsAlign, r)
+		}
+	}
+
+	for _, out := range []struct {
+		name string
+		err  error
+	}{
+		{*prefix + "_a.csv", knobloch.WriteTableA(*prefix+"_a.csv", rowsA)},
+		{*prefix + "_b.csv", knobloch.WriteTableB(*prefix+"_b.csv", rowsB)},
+		{*prefix + "_c_aligned.csv", knobloch.WriteTableC(*prefix+"_c_aligned.csv", aligned)},
+		{*prefix + "_c_inverted.csv", knobloch.WriteTableC(*prefix+"_c_inverted.csv", inverted)},
+		{*prefix + "_c_inverted_vs_aligned.csv", knobloch.WriteTableC(*prefix+"_c_inverted_vs_aligned.csv", invVsAlign)},
+	} {
+		if out.err != nil {
+			log.Fatal(out.err)
+		}
+
+		log.Printf("wrote %s", out.name)
+	}
+}