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
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
|
// fedisearch-spike is a vertical slice that proves the inference stack works
// before any application code is written. It downloads a sentence-transformer
// in ONNX format, runs it through GoMLX's pure-Go backend, and checks the
// resulting embedding against a known-good vector.
//
// It deliberately does NOT use github.com/knights-analytics/hugot. Hugot is a
// pleasant wrapper, but it costs 37 extra packages (viant/afs for cloud storage
// we never touch, x/crypto via its downloader) for ~400 lines of glue that we
// can own outright. This file is that glue, in its smallest honest form.
//
// Two things are being proven here, and only these two:
//
// 1. ONNX + GoMLX + the hftokenizer produce embeddings that match the
// reference implementation, with no cgo, no ONNX Runtime, no rust
// tokenizer.
// 2. The dependency closure (111 packages, notably google.golang.org/protobuf
// with its //go:embed) can be built by depot.nix.buildGo.
//
// Everything else — SQLite, ActivityPub ingest, images, the web UI — is
// deliberately absent. Once phase 2 lands a structured embedder, this file
// becomes a regression test rather than a program.
package main
import (
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"github.com/gomlx/compute"
// Imported for its init(), which calls compute.Register("go", ...).
// This is load-bearing: compute.NewWithConfig panics rather than returning
// an error when no backend is registered, so dropping this import fails at
// runtime, not at compile time.
_ "github.com/gomlx/compute/gobackend"
"github.com/gomlx/go-huggingface/tokenizers/api"
"github.com/gomlx/go-huggingface/tokenizers/hftokenizer"
"github.com/gomlx/gomlx/core/graph"
"github.com/gomlx/gomlx/core/tensors"
mlmodel "github.com/gomlx/gomlx/ml/model"
"github.com/gomlx/onnx-gomlx/onnx"
"github.com/gomlx/onnx-gomlx/onnx/parser"
)
// ---------------------------------------------------------------------------
// Model definition
// ---------------------------------------------------------------------------
// modelRepo is the HuggingFace repo we pull from. KnightsAnalytics' mirror is
// used rather than sentence-transformers' own because it keeps model.onnx at
// the repo root, and because it is the exact artefact the reference vectors in
// expectedRobertSmith were produced from.
const modelRepo = "KnightsAnalytics/all-MiniLM-L6-v2"
// maxSequenceLength truncates the tokenizer.
//
// Note this is 256, not the 512 that config.json's max_position_embeddings
// advertises. sentence-transformers ships a separate sentence_bert_config.json
// declaring max_seq_length: 256, and that is the length the model was actually
// trained and evaluated at. Using 512 would silently diverge from the Python
// implementation for long inputs.
const maxSequenceLength = 256
// modelFiles are fetched relative to the repo root.
//
// vocab.txt is not needed (tokenizer.json is self-contained), and neither is
// config.json: the sequence length is pinned above and the embedding dimension
// is read off the output tensor rather than trusted from metadata.
var modelFiles = []string{
"model.onnx",
"tokenizer.json",
}
// ---------------------------------------------------------------------------
// Reference vector
// ---------------------------------------------------------------------------
// expectedRobertSmith is the mean-pooled, *unnormalized* embedding of the
// string "robert smith", taken from hugot's testcases/embedded/vectors.json
// ("test1output"). Its L2 norm is ~7.691, which is how you can tell it is the
// raw pooled output rather than a normalized one.
//
// Only the first 16 of 384 dimensions are checked. That is ample: any error in
// tokenization, tensor layout, or pooling perturbs every dimension, so a
// mismatch shows up immediately in the first few.
var expectedRobertSmith = []float32{
-0.5136888, 0.32881597, -0.77382463, 0.21413505,
-0.21415548, 0.30707812, 0.69023287, -0.29770786,
0.27898467, 0.46517533, -0.37616116, 0.25144404,
0.4015117, -0.11901632, -0.39918411, 0.57255781,
}
// tolerance for the float comparison. The go backend accumulates in float32 and
// may reorder reductions relative to ONNX Runtime, so bit-identical output is
// not expected; 1e-4 is far tighter than any semantically meaningful drift
// while still catching real bugs.
const tolerance = 1e-4
// ---------------------------------------------------------------------------
// Model download
// ---------------------------------------------------------------------------
// ensureModel downloads the model into the user's cache directory if it is not
// already there, and returns the directory holding it.
//
// go-huggingface ships a `hub` package that does this, but it is deliberately
// not used: it drags in a downloader, file locking and uuid generation for what
// is three HTTP GETs, and phase 1 needs to record checksums and revisions in
// SQLite anyway, which means owning this path regardless.
func ensureModel() (string, error) {
cacheRoot, err := os.UserCacheDir()
if err != nil {
return "", fmt.Errorf("locating cache dir: %w", err)
}
modelDir := filepath.Join(cacheRoot, "fedisearch", "models", filepath.Base(modelRepo))
if err := os.MkdirAll(modelDir, 0o755); err != nil {
return "", fmt.Errorf("creating %s: %w", modelDir, err)
}
for _, name := range modelFiles {
dest := filepath.Join(modelDir, name)
if fi, err := os.Stat(dest); err == nil && fi.Size() > 0 {
continue
}
url := fmt.Sprintf("https://huggingface.co/%s/resolve/main/%s", modelRepo, name)
fmt.Printf("downloading %s ...\n", name)
if err := downloadFile(url, dest); err != nil {
return "", fmt.Errorf("downloading %s: %w", name, err)
}
}
return modelDir, nil
}
// downloadFile fetches url into dest, writing to a temporary file first so that
// an interrupted download cannot be mistaken for a complete one on the next run
// (which is what the size check in ensureModel would otherwise do).
func downloadFile(url, dest string) error {
resp, err := http.Get(url)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("GET %s: unexpected status %s", url, resp.Status)
}
tmp := dest + ".partial"
f, err := os.Create(tmp)
if err != nil {
return err
}
if _, err := io.Copy(f, resp.Body); err != nil {
f.Close()
os.Remove(tmp)
return err
}
if err := f.Close(); err != nil {
os.Remove(tmp)
return err
}
return os.Rename(tmp, dest)
}
// ---------------------------------------------------------------------------
// Embedder
// ---------------------------------------------------------------------------
// embedder holds a loaded ONNX model, its GoMLX executor and the matching
// tokenizer. It is not safe for concurrent use.
type embedder struct {
onnxModel onnx.Model
store *mlmodel.Store
backend compute.Backend
exec *mlmodel.Exec
tokenizer *hftokenizer.Tokenizer
inputNames []string
outputNames []string
// needsTypeIDs and needsAttentionMask record which optional inputs this
// particular model declares. BERT-family models take all three; others may
// not, and feeding an input the graph does not expect is an error.
needsTypeIDs bool
needsAttentionMask bool
}
func newEmbedder(modelDir string) (*embedder, error) {
onnxModel, err := parser.ParseFile(filepath.Join(modelDir, "model.onnx"))
if err != nil {
return nil, fmt.Errorf("parsing model.onnx: %w", err)
}
// Variables (the model weights) live in a Store; the graph is built against
// a Scope within it. This is GoMLX's equivalent of loading the weights.
store := mlmodel.NewStore()
scope := store.RootScope()
if err := onnxModel.VariablesToScope(scope); err != nil {
return nil, fmt.Errorf("loading weights into scope: %w", err)
}
inputNames, _ := onnxModel.Inputs()
outputNames, _ := onnxModel.Outputs()
// "go" selects the pure-Go backend registered by the blank import above.
// No bucketing is configured: the Go backend does not JIT-compile per
// shape, so padding inputs to bucket sizes would be pure waste.
backend, err := compute.NewWithConfig("go")
if err != nil {
return nil, fmt.Errorf("creating go backend: %w", err)
}
// callFunc maps the positional tensors handed to Exec back onto the ONNX
// graph's named inputs, in the order reported by onnxModel.Inputs().
callFunc := func(scope *mlmodel.Scope, inputs []*graph.Node) []*graph.Node {
named := make(map[string]*graph.Node, len(inputs))
for i, name := range inputNames {
named[name] = inputs[i]
}
return onnxModel.CallGraph(scope, inputs[0].Graph(), named, outputNames...)
}
exec, err := mlmodel.NewExec(backend, store, callFunc)
if err != nil {
return nil, fmt.Errorf("creating executor: %w", err)
}
e := &embedder{
onnxModel: onnxModel,
store: store,
backend: backend,
exec: exec,
inputNames: inputNames,
outputNames: outputNames,
}
for _, name := range inputNames {
switch name {
case "input_ids":
case "token_type_ids":
e.needsTypeIDs = true
case "attention_mask":
e.needsAttentionMask = true
default:
return nil, fmt.Errorf("model declares unsupported input %q", name)
}
}
tokenizerJSON, err := os.ReadFile(filepath.Join(modelDir, "tokenizer.json"))
if err != nil {
return nil, fmt.Errorf("reading tokenizer.json: %w", err)
}
tk, err := hftokenizer.NewFromContent(nil, tokenizerJSON)
if err != nil {
return nil, fmt.Errorf("parsing tokenizer.json: %w", err)
}
if err := tk.With(api.EncodeOptions{
AddSpecialTokens: true,
MaxLen: maxSequenceLength,
}); err != nil {
return nil, fmt.Errorf("configuring tokenizer: %w", err)
}
e.tokenizer = tk
return e, nil
}
func (e *embedder) close() {
if e.exec != nil {
e.exec.Finalize()
}
if e.store != nil {
e.store.Finalize()
}
if e.backend != nil {
e.backend.Finalize()
}
if e.onnxModel != nil {
_ = e.onnxModel.Close()
}
}
// embed runs a single string through the model and returns its mean-pooled
// sentence embedding.
//
// Batch size is fixed at one. Batching is a phase 2 concern; doing it here
// would mean handling ragged sequence lengths and padding masks for no gain in
// what this file sets out to prove.
func (e *embedder) embed(text string) ([]float32, error) {
flat, seqLen, hidden, err := e.forward(text)
if err != nil {
return nil, err
}
return meanPool(flat, seqLen, hidden), nil
}
// forward runs the model and returns the raw last_hidden_state as a flat
// array, together with its two dimensions: one row of `hidden` floats per
// token, `seqLen` rows.
//
// Split out from embed so that callers wanting the per-token vectors, rather
// than the pooled sentence vector, do not have to duplicate the tensor
// plumbing. chunkdemo.go uses it to show what pooling discards.
func (e *embedder) forward(text string) (flat []float32, seqLen, hidden int, err error) {
tokenIDs := e.tokenizer.Encode(text)
if len(tokenIDs) == 0 {
return nil, 0, 0, fmt.Errorf("tokenizer produced no tokens")
}
seqLen = len(tokenIDs)
// Every input is [batch=1, seqLen] int64. Because there is no padding in a
// batch of one, the attention mask is all ones and the token type IDs are
// all zeros — but they must still be supplied, as the graph declares them.
inputs := make([]*tensors.Tensor, len(e.inputNames))
for i, name := range e.inputNames {
data := make([]int64, seqLen)
switch name {
case "input_ids":
for j, id := range tokenIDs {
data[j] = int64(id)
}
case "attention_mask":
for j := range data {
data[j] = 1
}
case "token_type_ids":
// leave as zeros: single-segment input
}
inputs[i] = tensors.FromFlatDataAndDimensions(data, 1, seqLen)
}
defer func() {
for _, t := range inputs {
_ = t.FinalizeAll()
}
}()
outputs, err := e.exec.Exec(inputs)
if err != nil {
return nil, 0, 0, fmt.Errorf("running model: %w", err)
}
defer func() {
for _, t := range outputs {
_ = t.FinalizeAll()
}
}()
if len(outputs) == 0 {
return nil, 0, 0, fmt.Errorf("model returned no outputs")
}
// The first output is last_hidden_state, shaped [1, seqLen, hidden].
if err := tensors.ConstFlatData(outputs[0], func(data []float32) {
flat = append([]float32(nil), data...)
}); err != nil {
return nil, 0, 0, fmt.Errorf("reading output tensor: %w", err)
}
if len(flat)%seqLen != 0 {
return nil, 0, 0, fmt.Errorf("output length %d is not divisible by sequence length %d", len(flat), seqLen)
}
return flat, seqLen, len(flat) / seqLen, nil
}
// meanPool averages token embeddings into a single sentence vector.
//
// With a batch of one there is no padding, so every token participates. (Once
// batching arrives in phase 2 this needs an attention mask, otherwise padding
// tokens drag the mean toward zero — the classic way to get embeddings that
// look plausible but retrieve badly.)
//
// The result is intentionally left unnormalized, to match the reference vector.
func meanPool(flat []float32, seqLen, hidden int) []float32 {
out := make([]float32, hidden)
for t := range seqLen {
row := flat[t*hidden : (t+1)*hidden]
for i, v := range row {
out[i] += v
}
}
for i := range out {
out[i] /= float32(seqLen)
}
return out
}
|