summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--research/knobloch/cmd/train/main.go39
-rw-r--r--research/knobloch/cmd/train/serialize.go4
-rw-r--r--research/knobloch/frequency.go864
-rw-r--r--research/knobloch/morphemes.go127
-rw-r--r--tokenizer/bpe/utility.go4
5 files changed, 1017 insertions, 21 deletions
diff --git a/research/knobloch/cmd/train/main.go b/research/knobloch/cmd/train/main.go
index 72f7e00..ebe8c54 100644
--- a/research/knobloch/cmd/train/main.go
+++ b/research/knobloch/cmd/train/main.go
@@ -10,25 +10,24 @@ import (
"github.com/jonasknobloch/mbpe"
"go.jknobloc.com/x/dataset"
- "go.jknobloc.com/x/profile"
"go.jknobloc.com/x/shelf"
"go.jknobloc.com/x/tokenizer/bpe"
"go.jknobloc.com/x/tokenizer/bpe/split"
)
func main() {
- stop := profile.CPU()
+ // stop := profile.CPU()
- // train()
+ train()
// serialize()
- profile.Mem()
+ // profile.Mem()
- stop()
+ // stop()
}
func train() {
- out := shelf.Abs("results/knobloch/minipile")
+ out := shelf.Abs("results/knobloch/fineweb_verify")
if err := os.MkdirAll(out, os.ModePerm); err != nil {
log.Fatal(err)
@@ -44,7 +43,7 @@ func train() {
return m
}()
- // mbpe.InvertWeightFunction = true
+ mbpe.InvertWeightFunction = true
// mbpe.UseSimpleClashes = true
newTrainer := func(segmenter mbpe.Segmenter, alpha float64) *mbpe.MBPETrainer {
@@ -58,24 +57,24 @@ func train() {
b.SetMatcher(split.NewFSA())
- return mbpe.NewMBPETrainer(b, segmenter, alpha, mbpe.NewMBPE(), 1<<17, alphabet)
+ return mbpe.NewMBPETrainer(b, segmenter, alpha, mbpe.NewMBPE(), 1<<19, alphabet)
}
configs := []struct {
float64
string
}{
- {0.0, "m000_minipile_v2"},
- {0.1, "m010_minipile_v2"},
- {0.2, "m020_minipile_v2"},
- {0.3, "m030_minipile_v2"},
- {0.4, "m040_minipile_v2"},
- {0.5, "m050_minipile_v2"},
- {0.6, "m060_minipile_v2"},
- {0.7, "m070_minipile_v2"},
- {0.8, "m080_minipile_v2"},
- {0.9, "m090_minipile_v2"},
- {1.0, "m100_minipile_v2"},
+ {0.0, "mi000_fineweb"},
+ {0.1, "mi010_fineweb"},
+ {0.2, "mi020_fineweb"},
+ {0.3, "mi030_fineweb"},
+ {0.4, "mi040_fineweb"},
+ {0.5, "mi050_fineweb"},
+ {0.6, "mi060_fineweb"},
+ {0.7, "mi070_fineweb"},
+ {0.8, "mi080_fineweb"},
+ {0.9, "mi090_fineweb"},
+ {1.0, "mi100_fineweb"},
}
for i, c := range configs {
@@ -86,7 +85,7 @@ func train() {
if dictErr := t.LoadDict(dict); dictErr != nil {
var reader dataset.Reader
- if r, err := dataset.NewParquetReader(shelf.Abs("data/minipile/train")); err != nil {
+ if r, err := dataset.NewParquetReader(shelf.Abs("data/fineweb-edu/sample-10BT/train")); err != nil {
log.Fatal(err)
} else {
reader = r
diff --git a/research/knobloch/cmd/train/serialize.go b/research/knobloch/cmd/train/serialize.go
index 9fbf968..5536450 100644
--- a/research/knobloch/cmd/train/serialize.go
+++ b/research/knobloch/cmd/train/serialize.go
@@ -23,7 +23,9 @@ func serialize() {
paths = ps
}
- steps := []int{100512, 50256, 32768, 16384, 8192}
+ // steps := []int{100512, 50256, 32768, 16384, 8192}
+ // steps := []int{131072}
+ steps := []int{256}
outRoot := shelf.Abs("tokenizers")
diff --git a/research/knobloch/frequency.go b/research/knobloch/frequency.go
new file mode 100644
index 0000000..58638c8
--- /dev/null
+++ b/research/knobloch/frequency.go
@@ -0,0 +1,864 @@
+package knobloch
+
+import (
+ "fmt"
+ "image/color"
+ "log"
+ "slices"
+ "sync/atomic"
+ "time"
+
+ "github.com/jonasknobloch/mbpe"
+ "go.jknobloc.com/x/tokenizer/bpe"
+ "go.jknobloc.com/x/tui"
+ "gonum.org/v1/plot"
+ "gonum.org/v1/plot/plotter"
+ "gonum.org/v1/plot/vg"
+ "gonum.org/v1/plot/vg/draw"
+)
+
+// TODO we could color code intermediate / terminal tokens
+// TODO some mere layer like color coding would be even cooler
+// TODO we could even color code morpheme approximations?
+// TODO why not simply count runes? + mark terminal?
+
+// NRGBA; color.RGBA is alpha-premultiplied and mangles the blend
+var (
+ colorOther = color.NRGBA{R: 108, G: 126, B: 179, A: 96}
+ colorMorpheme = color.NRGBA{R: 214, G: 96, B: 77, A: 96}
+)
+
+func opaque(c color.NRGBA) color.NRGBA {
+ c.A = 255
+
+ return c
+}
+
+type NoPreTok struct{}
+
+func (p *NoPreTok) PreTokenize(s string) []string {
+ return []string{s}
+}
+
+func TokenFrequencies(t *bpe.Tokenizer, name, suffix string) error {
+ dict := mbpe.NewDict()
+
+ if err := dict.Load(name); err != nil {
+ return err
+ }
+
+ m := bpe.MBPE(t)
+
+ m.SetPreTokenizer(&NoPreTok{})
+
+ r := make([]int, len(bpe.Vocab(t)))
+
+ items := dict.Items()
+
+ pb := tui.NewProgressBar("Chunks", 20, len(items), time.Now())
+
+ var processed atomic.Int64
+
+ pb.Start(1*time.Second, func() int {
+ return int(processed.Load())
+ })
+
+ for _, v := range items {
+ n := v.N()
+
+ for _, id := range t.Encode(v.Src()) {
+ r[id] += n
+ }
+
+ processed.Add(1)
+ }
+
+ if err := plotAll(r, nil, t, suffix); err != nil {
+ return err
+ }
+
+ if len(SharedVocabs) > 0 {
+ shared, err := sharedMask(len(r), t, true)
+
+ if err != nil {
+ return err
+ }
+
+ if err := plotAll(r, shared, t, suffix+"_shared"); err != nil {
+ return err
+ }
+
+ unshared, err := sharedMask(len(r), t, false)
+
+ if err != nil {
+ return err
+ }
+
+ if err := plotAll(r, unshared, t, suffix+"_unshared"); err != nil {
+ return err
+ }
+
+ if err := plotMorphemeFractionOverview(r, shared, unshared, t, plotName("morpheme_fraction_overview", suffix)); err != nil {
+ return err
+ }
+
+ if err := plotMorphemeContribution(r, shared, unshared, t, plotName("morpheme_contribution", suffix)); err != nil {
+ return err
+ }
+
+ if err := plotMorphemeFractionLesci(r, shared, unshared, t, plotName("morpheme_fraction_lesci", suffix)); err != nil {
+ return err
+ }
+
+ if err := plotMorphemeFractionByIDOverview(r, shared, unshared, t, plotName("morpheme_fraction_by_id_overview", suffix)); err != nil {
+ return err
+ }
+ }
+
+ stats, err := NewFrequencyStats(suffix, r, t)
+
+ if err != nil {
+ return err
+ }
+
+ log.Printf("%s: %d/%d used tokens are morphemes (%.1f%% of types, %.1f%% of tokens), median %d vs %d, %d of top 1000, rank-id rho %.4f",
+ stats.Model, stats.MorphemeTypes, stats.Used, stats.TypeShare()*100, stats.TokenShare()*100,
+ stats.MorphemeMedian, stats.OtherMedian, stats.MorphemeTop1000, stats.RankID)
+
+ return AppendFrequencyStats("frequency_stats.csv", stats)
+}
+
+// plotAll draws every figure for one tokenizer. A nil keep covers the whole
+// vocabulary; otherwise the subset it selects is plotted. Note that the morpheme
+// fraction gets the unmasked counts, since it ranks over the full vocabulary
+// before hiding anything.
+func plotAll(r []int, keep []bool, t *bpe.Tokenizer, suffix string) error {
+ masked := r
+
+ if keep != nil {
+ masked = maskCounts(r, keep)
+ }
+
+ if err := plotFrequenciesByIDSplit(masked, t, plotName("frequencies_by_id_split", suffix)); err != nil {
+ return err
+ }
+
+ // two takes on the rank axis: mixed ranks the whole vocabulary and masks
+ // after, so subsets stay on the full rank space, while split ranks each
+ // group alone so the two separate
+ if err := plotFrequenciesMixed(r, keep, t, plotName("frequencies_mixed", suffix)); err != nil {
+ return err
+ }
+
+ if err := plotFrequenciesSplit(masked, t, plotName("frequencies_split", suffix)); err != nil {
+ return err
+ }
+
+ return plotMorphemeFraction(r, keep, t, plotName("morpheme_fraction", suffix))
+}
+
+func plotName(base, suffix string) string {
+ if suffix == "" {
+ return base + ".png"
+ }
+
+ return base + "_" + suffix + ".png"
+}
+
+func plotFrequenciesByIDSplit(m []int, t *bpe.Tokenizer, out string) error {
+ a, b, err := splitScatter(pointsByID(m), t)
+
+ if err != nil {
+ return err
+ }
+
+ return renderFrequenciesSplit(a, b, "Token ID", len(m), out)
+}
+
+// plotFrequenciesSplit ranks tokens by frequency rather than id. Each group is
+// ranked within itself: a shared ranking puts both on one curve, where the
+// groups overplot each other and no composition is visible.
+func plotFrequenciesSplit(m []int, t *bpe.Tokenizer, out string) error {
+ morph, err := isMorpheme(t)
+
+ if err != nil {
+ return err
+ }
+
+ var a, b []int
+
+ for id, v := range m {
+ if v <= 0 {
+ continue
+ }
+
+ if morph(id) {
+ a = append(a, v)
+ } else {
+ b = append(b, v)
+ }
+ }
+
+ return renderFrequenciesSplit(rankPoints(a), rankPoints(b), "Token frequency rank", len(m), out)
+}
+
+// plotFrequenciesMixed ranks every token on one shared scale and colors each
+// point by group, so composition is read along a single curve rather than from
+// the two separate curves plotFrequenciesSplit draws.
+//
+// Ranking covers the whole vocabulary and keep only hides points afterwards, so
+// a subset keeps the rank positions it holds in the full figure. Ranking within
+// the subset instead would stretch it back across the axis; that is what
+// plotFrequenciesSplit does, deliberately, as the other half of the pair.
+func plotFrequenciesMixed(m []int, keep []bool, t *bpe.Tokenizer, out string) error {
+ morph, err := isMorpheme(t)
+
+ if err != nil {
+ return err
+ }
+
+ ids := rankedIDs(m)
+
+ pts := make(plotter.XYs, 0, len(ids))
+ flags := make([]bool, 0, len(ids))
+
+ for i, id := range ids {
+ if keep != nil && !keep[id] {
+ continue
+ }
+
+ pts = append(pts, plotter.XY{X: float64(i + 1), Y: float64(m[id])})
+ flags = append(flags, morph(id))
+ }
+
+ if len(pts) == 0 {
+ return fmt.Errorf("no tokens selected")
+ }
+
+ return renderFrequenciesMixed(pts, flags, "Token frequency rank", len(m), out)
+}
+
+// rankWindow is wide enough to smooth the per-token noise but still leaves a few
+// hundred points across a typical vocabulary.
+const rankWindow = 250
+
+// plotMorphemeFraction plots the morpheme share of a window over the frequency
+// ranking. Encoding composition as position rather than hue keeps it readable,
+// which renderFrequenciesMixed cannot manage on a single curve.
+//
+// Ranking covers the whole vocabulary and keep only hides points afterwards, so
+// a subset lands on the same x positions it occupies in the full figure. Ranking
+// within the subset instead would stretch it back across the axis and make the
+// subsets impossible to compare.
+func plotMorphemeFraction(m []int, keep []bool, t *bpe.Tokenizer, out string) error {
+ morph, err := isMorpheme(t)
+
+ if err != nil {
+ return err
+ }
+
+ ids := rankedIDs(m)
+
+ if len(ids) == 0 {
+ return fmt.Errorf("no used tokens")
+ }
+
+ segments, overall, shown := morphemeFractionCurve(ids, keep, morph)
+
+ if shown == 0 {
+ return fmt.Errorf("no tokens selected")
+ }
+
+ return renderMorphemeFraction([]fractionSeries{{
+ segments: segments,
+ overall: overall,
+ ref: true,
+ color: opaque(colorMorpheme),
+ label: "Window share",
+ }}, "Token frequency rank", len(m), out)
+}
+
+// plotMorphemeFractionOverview overlays the full vocabulary and both sides of the
+// intersection. The shared ranking makes the three curves directly comparable,
+// so their overall shares go in the legend instead of as reference lines.
+func plotMorphemeFractionOverview(m []int, shared, unshared []bool, t *bpe.Tokenizer, out string) error {
+ morph, err := isMorpheme(t)
+
+ if err != nil {
+ return err
+ }
+
+ ids := rankedIDs(m)
+
+ if len(ids) == 0 {
+ return fmt.Errorf("no used tokens")
+ }
+
+ var series []fractionSeries
+
+ for _, c := range []struct {
+ keep []bool
+ color color.NRGBA
+ label string
+ }{
+ {nil, color.NRGBA{R: 90, G: 90, B: 90, A: 255}, "All"},
+ {shared, opaque(colorOther), "Shared"},
+ {unshared, opaque(colorMorpheme), "Unshared"},
+ } {
+ segments, overall, shown := morphemeFractionCurve(ids, c.keep, morph)
+
+ if shown == 0 {
+ continue
+ }
+
+ series = append(series, fractionSeries{
+ segments: segments,
+ overall: overall,
+ color: c.color,
+ label: fmt.Sprintf("%s (%.1f%%)", c.label, overall),
+ })
+ }
+
+ if len(series) == 0 {
+ return fmt.Errorf("no tokens selected")
+ }
+
+ return renderMorphemeFraction(series, "Token frequency rank", len(m), out)
+}
+
+// plotMorphemeContribution decomposes the overall morpheme share into the part
+// each side of the intersection supplies. Every curve is divided by all tokens
+// in the bin rather than by its own subset, so the two stack to exactly the
+// combined curve, which the rate view in plotMorphemeFractionOverview does not.
+//
+// The denominator also being the full bin removes the small-sample noise that
+// makes the rate view ragged wherever a subset is thin, and leaves no empty bins
+// to break the line. The cost is that a low value here mixes "morpheme-poor"
+// with "barely present", which the rate view keeps apart.
+func plotMorphemeContribution(m []int, shared, unshared []bool, t *bpe.Tokenizer, out string) error {
+ morph, err := isMorpheme(t)
+
+ if err != nil {
+ return err
+ }
+
+ ids := rankedIDs(m)
+
+ if len(ids) == 0 {
+ return fmt.Errorf("no used tokens")
+ }
+
+ sharedPts, sharedTotal := contributionCurve(ids, shared, morph)
+ _, unsharedTotal := contributionCurve(ids, unshared, morph)
+ combined, combinedTotal := contributionCurve(ids, nil, morph)
+
+ // the combined curve is filled to the baseline first, then the shared part is
+ // filled over it, so the salmon left showing is the unshared band. The same
+ // combined values are drawn again unfilled on top, since the total is a curve
+ // worth reading in its own right and not just the top edge of a stack.
+ return renderMorphemeContribution([]fractionSeries{
+ {
+ segments: []plotter.XYs{combined},
+ color: opaque(colorMorpheme),
+ fill: true,
+ label: fmt.Sprintf("Unshared (%.1f%%)", unsharedTotal),
+ },
+ {
+ segments: []plotter.XYs{sharedPts},
+ color: opaque(colorOther),
+ fill: true,
+ label: fmt.Sprintf("Shared (%.1f%%)", sharedTotal),
+ },
+ {
+ segments: []plotter.XYs{combined},
+ color: color.NRGBA{R: 90, G: 90, B: 90, A: 255},
+ label: fmt.Sprintf("Combined (%.1f%%)", combinedTotal),
+ },
+ }, combinedTotal, len(m), out)
+}
+
+// contributionCurve reports the morphemes selected by keep as a percentage of
+// every token in the bin, so curves for disjoint keeps add up.
+func contributionCurve(ids []int, keep []bool, morph func(int) bool) (plotter.XYs, float64) {
+ pts := make(plotter.XYs, 0, len(ids)/rankWindow+1)
+
+ morphemes := 0
+
+ for start := 0; start < len(ids); start += rankWindow {
+ end := min(start+rankWindow, len(ids))
+
+ count := 0
+
+ for _, id := range ids[start:end] {
+ if keep != nil && !keep[id] {
+ continue
+ }
+
+ if morph(id) {
+ count++
+ }
+ }
+
+ morphemes += count
+
+ pts = append(pts, plotter.XY{
+ X: float64(start + (end-start)/2),
+ Y: 100 * float64(count) / float64(end-start),
+ })
+ }
+
+ if len(ids) == 0 {
+ return pts, 0
+ }
+
+ return pts, 100 * float64(morphemes) / float64(len(ids))
+}
+
+func renderMorphemeContribution(series []fractionSeries, overall float64, vocab int, out string) error {
+ p := plot.New()
+
+ p.X.Label.Text = "Token frequency rank"
+ p.Y.Label.Text = "Share of all tokens in window (%)"
+
+ // pin to the vocabulary so figures stay comparable between tokenizers
+ p.X.Min = 0
+ p.X.Max = float64(vocab)
+
+ p.Y.Min = 0
+ p.Y.Max = 100
+
+ p.Add(plotter.NewGrid())
+
+ ref := plotter.NewFunction(func(float64) float64 {
+ return overall
+ })
+
+ ref.Color = color.NRGBA{R: 131, G: 131, B: 131, A: 255}
+ ref.Width = vg.Points(1)
+ ref.Dashes = []vg.Length{vg.Points(4), vg.Points(3)}
+
+ p.Add(ref)
+ p.Legend.Add("Combined mean", ref)
+
+ for _, s := range series {
+ for _, pts := range s.segments {
+ line, err := plotter.NewLine(pts)
+
+ if err != nil {
+ return err
+ }
+
+ line.Color = s.color
+ line.Width = vg.Points(1)
+
+ if s.fill {
+ // opaque: the shared fill has to occlude the combined fill
+ // underneath it, otherwise the combined area shows through and
+ // the visible red is the total rather than the unshared band
+ line.FillColor = opaque(s.color)
+ } else {
+ line.Width = vg.Points(1.5)
+ }
+
+ p.Add(line)
+ p.Legend.Add(s.label, line)
+ }
+ }
+
+ p.Legend.Top = true
+ p.Legend.Padding = vg.Points(4)
+
+ return p.Save(12*vg.Inch, 6*vg.Inch, out)
+}
+
+// idFractionCurve bins by token id instead of by rank, over the half-open id
+// range [lo, hi). Empty bins end the current segment rather than being
+// interpolated. A nil keep covers every token.
+func idFractionCurve(m []int, keep []bool, morph func(int) bool, lo, hi int) (segments []plotter.XYs, overall float64, shown int) {
+ var current plotter.XYs
+
+ morphemes := 0
+
+ for start := lo; start < hi; start += rankWindow {
+ end := min(start+rankWindow, hi)
+
+ count, total := 0, 0
+
+ for id := start; id < end; id++ {
+ if id < 0 || id >= len(m) || m[id] <= 0 || (keep != nil && !keep[id]) {
+ continue
+ }
+
+ total++
+
+ if morph(id) {
+ count++
+ }
+ }
+
+ if total == 0 {
+ if len(current) > 0 {
+ segments = append(segments, current)
+
+ current = nil
+ }
+
+ continue
+ }
+
+ shown += total
+ morphemes += count
+
+ current = append(current, plotter.XY{
+ X: float64(start + (end-start)/2),
+ Y: 100 * float64(count) / float64(total),
+ })
+ }
+
+ if len(current) > 0 {
+ segments = append(segments, current)
+ }
+
+ if shown > 0 {
+ overall = 100 * float64(morphemes) / float64(shown)
+ }
+
+ return segments, overall, shown
+}
+
+// plotMorphemeFractionByIDOverview is the token id counterpart of
+// plotMorphemeFractionOverview: same three series, binned by vocab position over
+// the whole vocabulary rather than by frequency rank, and with no lesci filter.
+func plotMorphemeFractionByIDOverview(m []int, shared, unshared []bool, t *bpe.Tokenizer, out string) error {
+ morph, err := isMorpheme(t)
+
+ if err != nil {
+ return err
+ }
+
+ var series []fractionSeries
+
+ for _, c := range []struct {
+ keep []bool
+ color color.NRGBA
+ label string
+ }{
+ {nil, color.NRGBA{R: 90, G: 90, B: 90, A: 255}, "All"},
+ {shared, opaque(colorOther), "Shared"},
+ {unshared, opaque(colorMorpheme), "Unshared"},
+ } {
+ segments, overall, shown := idFractionCurve(m, c.keep, morph, 0, len(m))
+
+ if shown == 0 {
+ continue
+ }
+
+ series = append(series, fractionSeries{
+ segments: segments,
+ overall: overall,
+ color: c.color,
+ label: fmt.Sprintf("%s (%.1f%%)", c.label, overall),
+ })
+ }
+
+ if len(series) == 0 {
+ return fmt.Errorf("no tokens selected")
+ }
+
+ return renderMorphemeFraction(series, "Token ID", len(m), out)
+}
+
+// rankedIDs lists the used vocab ids most frequent first.
+func rankedIDs(m []int) []int {
+ ids := make([]int, 0, len(m))
+
+ for id, v := range m {
+ if v > 0 {
+ ids = append(ids, id)
+ }
+ }
+
+ slices.SortFunc(ids, func(a, b int) int {
+ return m[b] - m[a]
+ })
+
+ return ids
+}
+
+// morphemeFractionCurve bins the ranking and reports the morpheme share of the
+// selected tokens per bin. Bins are fixed on the ranking it is given, so a bin
+// covers the same ranks for every keep; empty bins end the current segment
+// rather than letting the line interpolate across them.
+func morphemeFractionCurve(ids []int, keep []bool, morph func(int) bool) (segments []plotter.XYs, overall float64, shown int) {
+ var current plotter.XYs
+
+ morphemes := 0
+
+ for start := 0; start < len(ids); start += rankWindow {
+ end := min(start+rankWindow, len(ids))
+
+ count, total := 0, 0
+
+ for _, id := range ids[start:end] {
+ if keep != nil && !keep[id] {
+ continue
+ }
+
+ total++
+
+ if morph(id) {
+ count++
+ }
+ }
+
+ if total == 0 {
+ if len(current) > 0 {
+ segments = append(segments, current)
+
+ current = nil
+ }
+
+ continue
+ }
+
+ shown += total
+ morphemes += count
+
+ current = append(current, plotter.XY{
+ X: float64(start + (end-start)/2),
+ Y: 100 * float64(count) / float64(total),
+ })
+ }
+
+ if len(current) > 0 {
+ segments = append(segments, current)
+ }
+
+ if shown > 0 {
+ overall = 100 * float64(morphemes) / float64(shown)
+ }
+
+ return segments, overall, shown
+}
+
+type fractionSeries struct {
+ segments []plotter.XYs
+ overall float64
+ ref bool
+ fill bool
+ color color.NRGBA
+ label string
+}
+
+func renderMorphemeFraction(series []fractionSeries, labelX string, vocab int, out string) error {
+ p := plot.New()
+
+ p.X.Label.Text = labelX
+ p.Y.Label.Text = "Morphemes in window (%)"
+
+ // pin to the vocabulary so figures stay comparable between tokenizers
+ p.X.Min = 0
+ p.X.Max = float64(vocab)
+
+ p.Y.Min = 0
+ p.Y.Max = 100
+
+ p.Add(plotter.NewGrid())
+
+ for _, s := range series {
+ if s.ref {
+ ref := plotter.NewFunction(func(float64) float64 {
+ return s.overall
+ })
+
+ ref.Color = color.NRGBA{R: 131, G: 131, B: 131, A: 255}
+ ref.Width = vg.Points(1)
+ ref.Dashes = []vg.Length{vg.Points(4), vg.Points(3)}
+
+ p.Add(ref)
+ p.Legend.Add("Overall share", ref)
+ }
+
+ // one line per segment; only the first carries the legend entry
+ for i, pts := range s.segments {
+ line, err := plotter.NewLine(pts)
+
+ if err != nil {
+ return err
+ }
+
+ line.Color = s.color
+ line.Width = vg.Points(1.5)
+
+ p.Add(line)
+
+ if i == 0 {
+ p.Legend.Add(s.label, line)
+ }
+ }
+ }
+
+ p.Legend.Top = true
+ p.Legend.Padding = vg.Points(4)
+
+ return p.Save(12*vg.Inch, 6*vg.Inch, out)
+}
+
+func rankPoints(values []int) plotter.XYs {
+ slices.SortFunc(values, func(a, b int) int {
+ return b - a
+ })
+
+ pts := make(plotter.XYs, 0, len(values))
+
+ for i, v := range values {
+ pts = append(pts, plotter.XY{X: float64(i + 1), Y: float64(v)})
+ }
+
+ return pts
+}
+
+func pointsByID(m []int) plotter.XYs {
+ pts := make(plotter.XYs, 0, len(m))
+
+ for id, v := range m {
+ if v <= 0 {
+ continue
+ }
+
+ pts = append(pts, plotter.XY{X: float64(id), Y: float64(v)})
+ }
+
+ return pts
+}
+
+func splitScatter(pts plotter.XYs, t *bpe.Tokenizer) (a, b plotter.XYs, err error) {
+ morph, err := isMorpheme(t)
+
+ if err != nil {
+ return nil, nil, err
+ }
+
+ for _, pt := range pts {
+ if morph(int(pt.X)) {
+ a = append(a, pt)
+ } else {
+ b = append(b, pt)
+ }
+ }
+
+ return a, b, nil
+}
+
+func renderFrequenciesSplit(a, b plotter.XYs, labelX string, vocab int, out string) error {
+ p := plot.New()
+
+ p.X.Label.Text = labelX
+ p.Y.Label.Text = "Token frequency"
+
+ p.Y.Scale = plot.LogScale{}
+
+ p.Y.Tick.Marker = plot.LogTicks{Prec: -1}
+
+ // pin to the vocabulary so figures stay comparable between tokenizers
+ p.X.Min = 0
+ p.X.Max = float64(vocab)
+
+ p.Add(plotter.NewGrid())
+
+ addScatter := func(pts plotter.XYs, c color.NRGBA, label string) error {
+ s, err := plotter.NewScatter(pts)
+
+ if err != nil {
+ return err
+ }
+
+ s.Color = c
+ s.Radius = vg.Points(1)
+ s.Shape = draw.CircleGlyph{}
+
+ p.Add(s)
+
+ return addLegendKey(p, c, label)
+ }
+
+ if err := addScatter(b, colorOther, "Other"); err != nil {
+ return err
+ }
+
+ if err := addScatter(a, colorMorpheme, "Morpheme"); err != nil {
+ return err
+ }
+
+ p.Legend.Padding = vg.Points(4)
+
+ return p.Save(12*vg.Inch, 7*vg.Inch, out)
+}
+
+// renderFrequenciesMixed colors a single scatter per point. Overlaying two
+// plotters instead would draw one group entirely on top of the other, which
+// hides the composition wherever both share a rank range.
+func renderFrequenciesMixed(pts plotter.XYs, morph []bool, labelX string, vocab int, out string) error {
+ p := plot.New()
+
+ p.X.Label.Text = labelX
+ p.Y.Label.Text = "Token frequency"
+
+ p.Y.Scale = plot.LogScale{}
+
+ p.Y.Tick.Marker = plot.LogTicks{Prec: -1}
+
+ // pin to the vocabulary so figures stay comparable between tokenizers
+ p.X.Min = 0
+ p.X.Max = float64(vocab)
+
+ p.Add(plotter.NewGrid())
+
+ s, err := plotter.NewScatter(pts)
+
+ if err != nil {
+ return err
+ }
+
+ s.GlyphStyleFunc = func(i int) draw.GlyphStyle {
+ c := colorOther
+
+ if morph[i] {
+ c = colorMorpheme
+ }
+
+ return draw.GlyphStyle{Color: c, Radius: vg.Points(1), Shape: draw.CircleGlyph{}}
+ }
+
+ p.Add(s)
+
+ if err := addLegendKey(p, colorOther, "Other"); err != nil {
+ return err
+ }
+
+ if err := addLegendKey(p, colorMorpheme, "Morpheme"); err != nil {
+ return err
+ }
+
+ p.Legend.Padding = vg.Points(4)
+
+ return p.Save(12*vg.Inch, 7*vg.Inch, out)
+}
+
+// addLegendKey adds a solid, larger glyph; the plotted points are too small and
+// too faint to read as a legend key.
+func addLegendKey(p *plot.Plot, c color.NRGBA, label string) error {
+ thumb, err := plotter.NewScatter(plotter.XYs{})
+
+ if err != nil {
+ return err
+ }
+
+ thumb.Color = opaque(c)
+ thumb.Radius = vg.Points(3)
+ thumb.Shape = draw.CircleGlyph{}
+
+ p.Legend.Add(label, thumb)
+
+ return nil
+}
diff --git a/research/knobloch/morphemes.go b/research/knobloch/morphemes.go
new file mode 100644
index 0000000..226320c
--- /dev/null
+++ b/research/knobloch/morphemes.go
@@ -0,0 +1,127 @@
+package knobloch
+
+import (
+ "encoding/gob"
+ "io"
+ "os"
+ "sync"
+
+ "go.jknobloc.com/x/shelf"
+ "go.jknobloc.com/x/tokenizer/bpe"
+)
+
+type segmentsEntry struct {
+ Value []string
+ OK bool
+}
+
+// SegmentsPath holds the morphological segmentation that decides which vocab
+// tokens count as morphemes; resolved against the shelf root on first use.
+var SegmentsPath shelf.Item = "results/knobloch/minipile_19_fix/segments.gob"
+
+var morphemesCache struct {
+ once sync.Once
+ m map[string]int
+ err error
+}
+
+// morphemes caches the parsed segmentation; the gob is several hundred MB and
+// every plot and stats pass needs the same map.
+func morphemes() (map[string]int, error) {
+ morphemesCache.once.Do(func() {
+ morphemesCache.m, morphemesCache.err = Morphemes(shelf.Abs(SegmentsPath))
+ })
+
+ return morphemesCache.m, morphemesCache.err
+}
+
+// isMorpheme reports whether a vocab id spells out a known morpheme; the
+// returned closure holds the id-to-token map so callers build it once.
+func isMorpheme(t *bpe.Tokenizer) (func(id int) bool, error) {
+ m, err := morphemes()
+
+ if err != nil {
+ return nil, err
+ }
+
+ itoa := bpe.Itoa(t)
+
+ return func(id int) bool {
+ token, ok := itoa[int64(id)]
+
+ if !ok {
+ panic("unexpected token ID")
+ }
+
+ _, isMorph := m[token]
+
+ return isMorph
+ }, nil
+}
+
+func Morphemes(name string) (map[string]int, error) {
+ var file *os.File
+
+ if f, err := os.Open(name); err != nil {
+ return nil, err
+ } else {
+ file = f
+
+ defer file.Close()
+ }
+
+ var segments map[string]segmentsEntry
+
+ if m, err := loadSegments(file); err != nil {
+ return nil, err
+ } else {
+ segments = m
+ }
+
+ morphemes := make(map[string]int)
+
+ for _, v := range segments {
+ // the segmenter returns the unsegmented compound when it fails
+ if !v.OK || len(v.Value) == 1 {
+ continue
+ }
+
+ for _, s := range v.Value {
+ morphemes[s]++
+ }
+ }
+
+ return morphemes, nil
+
+ // r := make([]int64, 0)
+
+ // atoi := bpe.Atoi(t)
+ //
+ // for k, v := range morphemes {
+ // if v < 0 {
+ // continue
+ // }
+ //
+ // i, ok := atoi[k]
+ //
+ // if !ok {
+ // continue
+ // }
+ //
+ // r = append(r, i)
+ // }
+ //
+ // return r, nil
+}
+
+func loadSegments(r io.Reader) (map[string]segmentsEntry, error) {
+ m := make(map[string]segmentsEntry)
+
+ dec := gob.NewDecoder(r)
+
+ if err := dec.Decode(&m); err != nil {
+ return nil, err
+ }
+
+ return m, nil
+}
diff --git a/tokenizer/bpe/utility.go b/tokenizer/bpe/utility.go
index f8fcdbc..79e6faf 100644
--- a/tokenizer/bpe/utility.go
+++ b/tokenizer/bpe/utility.go
@@ -19,6 +19,10 @@ func NewTokenizerFromFiles(vocab, merges string, cfg Config) (*Tokenizer, error)
return NewTokenizer(tokenizer, cfg), nil
}
+func MBPE(t *Tokenizer) *mbpe.Tokenizer {
+ return t.mbpe
+}
+
func Vocab(t *Tokenizer) []string {
m, ok := t.mbpe.Model().(*mbpe.MBPE)
generated by cgit v1.2.3 (git 2.25.1) at 2026-10-06 23:41:43 +0000