1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
|
// Command tablea writes table A for a chosen set of tokenizers rather than the
// whole grid.
//
// Table A is a per-tokenizer statistic, so unlike tables B and C it needs no
// intersection and therefore no encode of the tokenizers that are not shown.
// Restricting the set cuts the work proportionally: three alignment levels at
// five vocabulary sizes is fifteen encodes instead of the full hundred and ten.
//
// The alignment column carries an " inv" suffix for the inverted family, which
// the full-grid command does not add.
package main
import (
"flag"
"fmt"
"log"
"runtime"
"strconv"
"strings"
"sync"
"github.com/jonasknobloch/mbpe"
"go.jknobloc.com/x/research/knobloch"
"go.jknobloc.com/x/shelf"
)
var vocabSizes = []int{8192, 16384, 32768, 50256, 100512}
// spec derives the alignment level from the tokenizer name, e.g. m050 or mi100.
func spec(size int, name string) (knobloch.TableSpec, error) {
inverted := strings.HasPrefix(name, "mi")
digits := strings.TrimPrefix(strings.TrimPrefix(name, "mi"), "m")
alpha, err := strconv.Atoi(digits)
if err != nil {
return knobloch.TableSpec{}, fmt.Errorf("%s: cannot read an alignment level from the name", name)
}
alignment := fmt.Sprintf("%.1f", float64(alpha)/100)
if inverted {
alignment += " inv"
}
return knobloch.TableSpec{
Name: name,
Alignment: alignment,
Inverted: inverted,
VocabSize: size,
Dir: shelf.Item(fmt.Sprintf("tokenizers/minipile/tokenizer_gpt2_%d_%s_minipile", size, name)),
}, nil
}
func main() {
dict := flag.String("dict", "results/knobloch/minipile/dict.txt", "shelf-relative dictionary")
names := flag.String("tokenizers", "m000,m050,m100", "comma separated tokenizer names")
sizes := flag.String("sizes", "", "comma separated vocabulary sizes, empty for all")
out := flag.String("out", "table_a_alphas.csv", "output CSV")
workers := flag.Int("workers", runtime.NumCPU(), "parallel encodes")
flag.Parse()
selected := vocabSizes
if *sizes != "" {
selected = nil
for _, s := range strings.Split(*sizes, ",") {
v, err := strconv.Atoi(strings.TrimSpace(s))
if err != nil {
log.Fatalf("bad vocabulary size %q", s)
}
selected = append(selected, v)
}
}
var specs []knobloch.TableSpec
for _, size := range selected {
for _, name := range strings.Split(*names, ",") {
s, err := spec(size, strings.TrimSpace(name))
if err != nil {
log.Fatal(err)
}
specs = append(specs, s)
}
}
d := mbpe.NewDict()
if err := d.Load(shelf.Abs(shelf.Item(*dict))); err != nil {
log.Fatal(err)
}
items := d.Items()
log.Printf("dictionary: %d pre-token types", len(items))
log.Printf("%d encodes: %s at %v", len(specs), *names, selected)
rows := make([]knobloch.TableARow, len(specs))
var wg sync.WaitGroup
queue := make(chan int)
for i := 0; i < *workers; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for idx := range queue {
e, err := knobloch.EncodeDict(specs[idx], items)
if err != nil {
log.Fatal(err)
}
rows[idx] = knobloch.BuildTableA(e)
log.Printf("%s at %d: encoded", specs[idx].Name, specs[idx].VocabSize)
}
}()
}
for i := range specs {
queue <- i
}
close(queue)
wg.Wait()
if err := knobloch.WriteTableA(*out, rows); err != nil {
log.Fatal(err)
}
log.Printf("wrote %s", *out)
}
|
