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)
			}
		}
	}
}