diff options
Diffstat (limited to 'research/knobloch/frequency.go')
| -rw-r--r-- | research/knobloch/frequency.go | 864 |
1 files changed, 864 insertions, 0 deletions
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 +} |
