WIPwip-frequency

1 files changed, 183 insertions, 0 deletions

diff --git a/research/knobloch/cmd/morphbands/main.go b/research/knobloch/cmd/morphbands/main.go
new file mode 100644
index 0000000..620b563
--- /dev/null
+++ b/research/knobloch/cmd/morphbands/main.go
@@ -0,0 +1,183 @@
+// Command morphbands reports morpheme content per band of vocabulary-id space,
+// counted two ways: by type and weighted by corpus frequency.
+//
+// The two disagree sharply. Morpheme types concentrate early -- around two
+// thirds of the first 2,048 sub-words -- and thin out into the tail, but the
+// tail is where most of the vocabulary is, so a type count says little about how
+// much text is actually covered by morphemic sub-words. Weighting each sub-word
+// by how often the tokenizer emits it over the corpus answers that instead.
+//
+// Frequencies come from encoding the dictionary, the same way the frequency
+// plots do, so this takes a few minutes per tokenizer and is the slow part.
+package main
+
+import (
+	"encoding/csv"
+	"flag"
+	"fmt"
+	"log"
+	"os"
+	"strconv"
+	"strings"
+
+	"github.com/jonasknobloch/mbpe"
+	"go.jknobloc.com/x/research/knobloch"
+	"go.jknobloc.com/x/shelf"
+	"go.jknobloc.com/x/tokenizer/bpe"
+)
+
+// frequencies counts how often each vocabulary id is emitted over the
+// dictionary, each word weighted by its corpus count.
+func frequencies(t *bpe.Tokenizer, dictPath string) ([]int, error) {
+	dict := mbpe.NewDict()
+
+	if err := dict.Load(dictPath); err != nil {
+		return nil, err
+	}
+
+	// the dictionary is already pre-tokenized, so re-splitting would corrupt it
+	m := bpe.MBPE(t)
+	m.SetPreTokenizer(&knobloch.NoPreTok{})
+
+	r := make([]int, len(bpe.Vocab(t)))
+
+	for _, v := range dict.Items() {
+		n := v.N()
+
+		for _, id := range t.Encode(v.Src()) {
+			r[id] += n
+		}
+	}
+
+	return r, nil
+}
+
+func parseBands(s string, size int) []int {
+	var edges []int
+
+	for _, f := range strings.Split(s, ",") {
+		v, err := strconv.Atoi(strings.TrimSpace(f))
+
+		if err != nil {
+			log.Fatalf("bad band edge %q: %v", f, err)
+		}
+
+		edges = append(edges, v)
+	}
+
+	if edges[0] != 0 {
+		edges = append([]int{0}, edges...)
+	}
+
+	if last := edges[len(edges)-1]; last < size {
+		edges = append(edges, size)
+	}
+
+	return edges
+}
+
+func share(part, whole int) string {
+	if whole == 0 {
+		return ""
+	}
+
+	return fmt.Sprintf("%.2f", 100*float64(part)/float64(whole))
+}
+
+func main() {
+	dirs := flag.String("dirs", "results/knobloch/minipile_19_ctrl/%s_minipile",
+		"tokenizer directory, %s is the alignment name")
+	dict := flag.String("dict", "results/knobloch/minipile_19_ctrl/dict.txt",
+		"pre-tokenized corpus dictionary")
+	list := flag.String("alignments", "m000,m050,m100,mi050,mi100",
+		"comma-separated alignments to report")
+	bands := flag.String("bands", "2048,8192,32768,100512,131072,262144",
+		"comma-separated upper edges of the id bands")
+	out := flag.String("out", "morpheme_bands.csv", "output CSV")
+
+	flag.Parse()
+
+	morph, err := knobloch.Morphemes(shelf.Abs(knobloch.SegmentsPath))
+
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	log.Printf("%d known morphemes", len(morph))
+
+	f, err := os.Create(*out)
+
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	defer f.Close()
+
+	w := csv.NewWriter(f)
+	defer w.Flush()
+
+	if err := w.Write([]string{"tokenizer", "band_lo", "band_hi", "size",
+		"morpheme_types", "type_share", "occurrences", "morpheme_occurrences",
+		"token_share"}); err != nil {
+		log.Fatal(err)
+	}
+
+	for _, n := range strings.Split(*list, ",") {
+		n = strings.TrimSpace(n)
+		dir := shelf.Item(fmt.Sprintf(*dirs, n))
+
+		tok, err := bpe.NewTokenizerFromFiles(
+			shelf.Abs(dir+"/vocab.json"), shelf.Abs(dir+"/merges.txt"),
+			bpe.Config{Recover: false})
+
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		vocab := bpe.Vocab(tok)
+
+		log.Printf("%s: encoding dictionary over %d sub-words", n, len(vocab))
+
+		r, err := frequencies(tok, shelf.Abs(shelf.Item(*dict)))
+
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		edges := parseBands(*bands, len(vocab))
+
+		log.Printf("%-6s %10s %9s %9s %11s %9s", n, "band", "types%", "tokens%", "occurrences", "morphs")
+
+		for i := 0; i+1 < len(edges); i++ {
+			lo, hi := edges[i], edges[i+1]
+
+			var types, morphTypes, occ, morphOcc int
+
+			for id := lo; id < hi && id < len(vocab); id++ {
+				types++
+				occ += r[id]
+
+				if _, ok := morph[vocab[id]]; ok {
+					morphTypes++
+					morphOcc += r[id]
+				}
+			}
+
+			row := []string{n, strconv.Itoa(lo), strconv.Itoa(hi), strconv.Itoa(types),
+				strconv.Itoa(morphTypes), share(morphTypes, types),
+				strconv.Itoa(occ), strconv.Itoa(morphOcc), share(morphOcc, occ)}
+
+			if err := w.Write(row); err != nil {
+				log.Fatal(err)
+			}
+
+			log.Printf("%-6s %10s %8s%% %8s%% %11d %9d", "",
+				fmt.Sprintf("%d-%d", lo, hi), share(morphTypes, types),
+				share(morphOcc, occ), occ, morphTypes)
+		}
+
+		w.Flush()
+	}
+
+	log.Printf("wrote %s", *out)
+}