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
145
146
147
148
149
150
151
152
|
package knobloch
import (
"encoding/json"
"fmt"
"os"
"sync"
"go.jknobloc.com/x/shelf"
"go.jknobloc.com/x/tokenizer/bpe"
)
// SharedVocabs lists the vocabularies intersected for the *_shared plots, which
// restrict every figure to tokens all of these models have in common. Set it to
// the family being swept so a model is only compared against its own kind; an
// empty list skips those plots.
var SharedVocabs = []shelf.Item{
"models/mbpe/minipile/gpt2_50256_m000_minipile/vocab.json",
"models/mbpe/minipile/gpt2_50256_m030_minipile/vocab.json",
"models/mbpe/minipile/gpt2_50256_m050_minipile/vocab.json",
"models/mbpe/minipile/gpt2_50256_m100_minipile/vocab.json",
}
var sharedCache struct {
mu sync.Mutex
m map[string]map[string]struct{}
}
// sharedTokens caches the intersection per vocabulary list, so a sweep that
// switches SharedVocabs between families gets each family's own intersection
// while still reading the vocabularies only once per family.
func sharedTokens() (map[string]struct{}, error) {
key := ""
for _, v := range SharedVocabs {
key += string(v) + "\n"
}
sharedCache.mu.Lock()
defer sharedCache.mu.Unlock()
if tokens, ok := sharedCache.m[key]; ok {
return tokens, nil
}
tokens, err := intersectVocabs(SharedVocabs)
if err != nil {
return nil, err
}
if sharedCache.m == nil {
sharedCache.m = make(map[string]map[string]struct{})
}
sharedCache.m[key] = tokens
return tokens, nil
}
func intersectVocabs(vocabs []shelf.Item) (map[string]struct{}, error) {
if len(vocabs) == 0 {
return nil, fmt.Errorf("no vocabularies to intersect")
}
var keep map[string]struct{}
for _, v := range vocabs {
tokens, err := loadVocab(shelf.Abs(v))
if err != nil {
return nil, err
}
if keep == nil {
keep = tokens
continue
}
for token := range keep {
if _, ok := tokens[token]; !ok {
delete(keep, token)
}
}
}
return keep, nil
}
func loadVocab(name string) (map[string]struct{}, error) {
file, err := os.Open(name)
if err != nil {
return nil, err
}
defer file.Close()
var m map[string]int64
if err := json.NewDecoder(file).Decode(&m); err != nil {
return nil, err
}
tokens := make(map[string]struct{}, len(m))
for token := range m {
tokens[token] = struct{}{}
}
return tokens, nil
}
// sharedMask selects the vocab ids on one side of the intersection: shared marks
// the tokens every vocabulary in SharedVocabs has, and its complement isolates
// what a listed model added on its own.
func sharedMask(n int, t *bpe.Tokenizer, shared bool) ([]bool, error) {
keep, err := sharedTokens()
if err != nil {
return nil, err
}
itoa := bpe.Itoa(t)
mask := make([]bool, n)
for id := range mask {
_, ok := keep[itoa[int64(id)]]
mask[id] = ok == shared
}
return mask, nil
}
// maskCounts zeroes the counts outside the mask. Every plot skips zero counts
// already, so the existing set of figures works unchanged while the axes stay
// pinned to the full vocabulary.
func maskCounts(r []int, keep []bool) []int {
masked := make([]int, len(r))
for id, v := range r {
if keep[id] {
masked[id] = v
}
}
return masked
}
|
