WIPwip-frequency

1 files changed, 58 insertions, 0 deletions

diff --git a/research/knobloch/cmd/plot/main.go b/research/knobloch/cmd/plot/main.go
new file mode 100644
index 0000000..e754f6f
--- /dev/null
+++ b/research/knobloch/cmd/plot/main.go
@@ -0,0 +1,58 @@
+package main
+
+import (
+	"fmt"
+	"log"
+
+	"go.jknobloc.com/x/research/knobloch"
+	"go.jknobloc.com/x/shelf"
+	"go.jknobloc.com/x/tokenizer/bpe"
+)
+
+const modelDir = "models/mbpe/minipile"
+
+// the two families are swept separately, so a model is only ever compared
+// against its own kind and the shared/unshared split never mixes m with mi
+var families = []string{"m", "mi"}
+
+func models(family string) []string {
+	var r []string
+
+	for alpha := 0; alpha <= 100; alpha += 10 {
+		r = append(r, fmt.Sprintf("gpt2_50256_%s%03d_minipile", family, alpha))
+	}
+
+	return r
+}
+
+func main() {
+	dict := shelf.Abs("results/knobloch/minipile/dict.txt")
+
+	cfg := bpe.Config{
+		Recover: false,
+	}
+
+	for _, family := range families {
+		names := models(family)
+
+		knobloch.SharedVocabs = nil
+
+		for _, model := range names {
+			knobloch.SharedVocabs = append(knobloch.SharedVocabs, shelf.Item(modelDir+"/"+model+"/vocab.json"))
+		}
+
+		for _, model := range names {
+			dir := shelf.Item(modelDir + "/" + model)
+
+			t, err := bpe.NewTokenizerFromFiles(shelf.Abs(dir+"/vocab.json"), shelf.Abs(dir+"/merges.txt"), cfg)
+
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			if err := knobloch.TokenFrequencies(t, dict, model); err != nil {
+				log.Fatal(err)
+			}
+		}
+	}
+}