1 files changed, 244 insertions, 0 deletions
diff --git a/research/knobloch/stats.go b/research/knobloch/stats.go
new file mode 100644
index 0000000..e462e1b
--- /dev/null
+++ b/research/knobloch/stats.go
@@ -0,0 +1,244 @@
+package knobloch
+
+import (
+	"encoding/csv"
+	"os"
+	"slices"
+	"strconv"
+
+	"go.jknobloc.com/x/tokenizer/bpe"
+	"gonum.org/v1/gonum/stat"
+)
+
+// FrequencyStats summarizes how much of a tokenizer's vocabulary, and how much
+// of its token mass, lines up with the morpheme inventory.
+type FrequencyStats struct {
+	Model string
+
+	// types are distinct vocabulary entries, occurrences are corpus counts
+	Vocab          int
+	Used           int
+	MorphemeTypes  int
+	MorphemeTokens int64
+	OtherTokens    int64
+
+	MorphemeMedian int
+	OtherMedian    int
+
+	// morphemes among the 1000 most frequent tokens
+	MorphemeTop1000 int
+
+	// RankID is Spearman's correlation between a token's vocabulary position and
+	// its corpus frequency rank. A strict frequency objective merges the most
+	// frequent pair first, so position and frequency rank move together and this
+	// sits near 1. The further a tokenizer departs from that objective, the more
+	// the two come apart, so it measures deviation from the frequency criterion
+	// independently of what the competing criterion happens to be.
+	RankID float64
+}
+
+func (s FrequencyStats) TypeShare() float64 {
+	if s.Used == 0 {
+		return 0
+	}
+
+	return float64(s.MorphemeTypes) / float64(s.Used)
+}
+
+func (s FrequencyStats) TokenShare() float64 {
+	total := s.MorphemeTokens + s.OtherTokens
+
+	if total == 0 {
+		return 0
+	}
+
+	return float64(s.MorphemeTokens) / float64(total)
+}
+
+func NewFrequencyStats(model string, r []int, t *bpe.Tokenizer) (FrequencyStats, error) {
+	morph, err := morphemes()
+
+	if err != nil {
+		return FrequencyStats{}, err
+	}
+
+	itoa := bpe.Itoa(t)
+
+	s := FrequencyStats{
+		Model: model,
+		Vocab: len(r),
+	}
+
+	var mFreq, oFreq []int
+
+	for id, f := range r {
+		if f <= 0 {
+			continue
+		}
+
+		s.Used++
+
+		if _, ok := morph[itoa[int64(id)]]; ok {
+			s.MorphemeTypes++
+			s.MorphemeTokens += int64(f)
+
+			mFreq = append(mFreq, f)
+		} else {
+			s.OtherTokens += int64(f)
+
+			oFreq = append(oFreq, f)
+		}
+	}
+
+	slices.Sort(mFreq)
+	slices.Sort(oFreq)
+
+	s.MorphemeMedian = median(mFreq)
+	s.OtherMedian = median(oFreq)
+
+	ids := make([]int, 0, s.Used)
+
+	for id, f := range r {
+		if f > 0 {
+			ids = append(ids, id)
+		}
+	}
+
+	slices.SortFunc(ids, func(a, b int) int {
+		return r[b] - r[a]
+	})
+
+	for _, id := range ids[:min(1000, len(ids))] {
+		if _, ok := morph[itoa[int64(id)]]; ok {
+			s.MorphemeTop1000++
+		}
+	}
+
+	s.RankID = rankIDCorrelation(r)
+
+	return s, nil
+}
+
+// rankIDCorrelation is Spearman's rho between vocabulary position and corpus
+// frequency rank over the used tokens. gonum has no Spearman, so the ranks are
+// built here and fed to Pearson, which is the same thing by definition.
+func rankIDCorrelation(r []int) float64 {
+	ids := make([]int, 0, len(r))
+
+	for id, f := range r {
+		if f > 0 {
+			ids = append(ids, id)
+		}
+	}
+
+	if len(ids) < 2 {
+		return 0
+	}
+
+	// ids is ascending and has no duplicates, so a token's id rank is its index
+	position := make([]float64, len(ids))
+
+	for i := range ids {
+		position[i] = float64(i + 1)
+	}
+
+	order := slices.Clone(ids)
+
+	slices.SortFunc(order, func(a, b int) int {
+		return r[b] - r[a]
+	})
+
+	// tied counts share the average of the ranks they span
+	ranks := make(map[int]float64, len(order))
+
+	for i := 0; i < len(order); {
+		j := i
+
+		for j+1 < len(order) && r[order[j+1]] == r[order[i]] {
+			j++
+		}
+
+		avg := float64(i+j+2) / 2
+
+		for k := i; k <= j; k++ {
+			ranks[order[k]] = avg
+		}
+
+		i = j + 1
+	}
+
+	frequency := make([]float64, len(ids))
+
+	for i, id := range ids {
+		frequency[i] = ranks[id]
+	}
+
+	return stat.Correlation(position, frequency, nil)
+}
+
+func median(v []int) int {
+	if len(v) == 0 {
+		return 0
+	}
+
+	return v[len(v)/2]
+}
+
+var frequencyStatsHeader = []string{
+	"model", "vocab", "used", "morpheme_types", "type_share",
+	"morpheme_tokens", "other_tokens", "token_share",
+	"morpheme_median", "other_median", "morpheme_top_1000", "rank_id_spearman",
+}
+
+func (s FrequencyStats) row() []string {
+	return []string{
+		s.Model,
+		strconv.Itoa(s.Vocab),
+		strconv.Itoa(s.Used),
+		strconv.Itoa(s.MorphemeTypes),
+		strconv.FormatFloat(s.TypeShare(), 'f', 4, 64),
+		strconv.FormatInt(s.MorphemeTokens, 10),
+		strconv.FormatInt(s.OtherTokens, 10),
+		strconv.FormatFloat(s.TokenShare(), 'f', 4, 64),
+		strconv.Itoa(s.MorphemeMedian),
+		strconv.Itoa(s.OtherMedian),
+		strconv.Itoa(s.MorphemeTop1000),
+		strconv.FormatFloat(s.RankID, 'f', 4, 64),
+	}
+}
+
+// AppendFrequencyStats adds one row per run so numbers for different tokenizers
+// accumulate in a single table.
+func AppendFrequencyStats(name string, s FrequencyStats) error {
+	_, err := os.Stat(name)
+
+	if err != nil && !os.IsNotExist(err) {
+		return err
+	}
+
+	fresh := os.IsNotExist(err)
+
+	file, err := os.OpenFile(name, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
+
+	if err != nil {
+		return err
+	}
+
+	defer file.Close()
+
+	w := csv.NewWriter(file)
+
+	if fresh {
+		if err := w.Write(frequencyStatsHeader); err != nil {
+			return err
+		}
+	}
+
+	if err := w.Write(s.row()); err != nil {
+		return err
+	}
+
+	w.Flush()
+
+	return w.Error()
+}