Something went wrong. Try again.
sentence embeddings in pure zig: bge-small with an HF-exact tokenizer and an SME matmul
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511//! BERT sentence encoder: the architecture, pooling and normalization come//! from the snapshot (spec.zig), the weights from model.safetensors.
const std = @import("std");const Allocator = std.mem.Allocator;const Io = std.Io;const ops = @import("ops.zig");const safetensors = @import("safetensors.zig");const spec_mod = @import("spec.zig");const Tokenizer = @import("tokenizer.zig").Tokenizer;const bf16 = @import("bf16.zig");const int8 = @import("int8.zig");
pub const Spec = spec_mod.Spec;pub const Pooling = spec_mod.Pooling;
pub const Precision = enum { f32, /// AVX-512 BF16 matmuls when the CPU has them and the kernel passes its /// self test; otherwise f32 bf16, /// bf16 numerics computed in plain Zig, for measuring bf16 accuracy on /// any machine; slow bf16_emulated, /// int8 weights and 12-bit activations on AVX2 when the kernel passes its /// self test; otherwise f32 int8, /// int8 numerics computed in plain Zig; slow int8_emulated, /// bf16 if this CPU runs it, else int8 if it runs that, else f32 fastest,};
/// Load options. Everything about the model itself is read from the snapshot.pub const Config = struct { precision: Precision = .f32,};
/// Optional per-stage time accounting. The caller supplies the clock (embedz/// has no Io or libc dependency), e.g. a thread-CPU clock.pub const Profile = struct { now: *const fn () u64, ns: std.EnumArray(Stage, std.atomic.Value(u64)) = .initFill(.init(0)), tokens: std.atomic.Value(u64) = .init(0),
pub const Stage = enum { tokenize, embeddings, qkv, attention, attn_out, ffn_in, gelu, ffn_out, layernorm, pool };
fn add(self: *Profile, stage: Stage, since: u64) u64 { const t = self.now(); _ = self.ns.getPtr(stage).fetchAdd(t -% since, .monotonic); return t; }};
const Layer = struct { qkv: ops.Linear, attn_out: ops.Linear, ln1_g: []f32, ln1_b: []f32, ffn_in: ops.Linear, ffn_out: ops.Linear, ln2_g: []f32, ln2_b: []f32,};
pub const Model = struct { spec: Spec, tokenizer: Tokenizer, word: []f32, /// position embeddings with token-type 0 already added pos: []f32, emb_ln_g: []f32, emb_ln_b: []f32, layers: []Layer, profile: ?*Profile = null, /// what `Config.precision` resolved to on this machine precision: Precision,
/// Load a sentence-transformers BERT snapshot directory: config.json, /// 1_Pooling/config.json, modules.json, sentence_bert_config.json, /// tokenizer.json, vocab.txt and model.safetensors. pub fn load(gpa: Allocator, io: Io, dir_path: []const u8, cfg: Config) !Model { var dir = try Io.Dir.cwd().openDir(io, dir_path, .{}); defer dir.close(io); return loadDir(gpa, io, dir, cfg); }
pub fn loadDir(gpa: Allocator, io: Io, dir: Io.Dir, cfg: Config) !Model { var spec = try spec_mod.read(gpa, io, dir);
const vocab_txt = try dir.readFileAlloc(io, "vocab.txt", gpa, .limited(16 << 20)); defer gpa.free(vocab_txt); var tokenizer = try Tokenizer.init(gpa, vocab_txt); errdefer tokenizer.deinit(gpa); if (tokenizer.vocab.count() != spec.vocab) return error.VocabSizeMismatch;
// mapped rather than read: the packed copy is the only heap-resident one, // so peak memory is ~1x the weights instead of ~2x const weights = try dir.openFile(io, "model.safetensors", .{}); defer weights.close(io); const len: usize = @intCast(try weights.length(io)); var map = try weights.createMemoryMap(io, .{ .len = len, .protection = .{ .read = true }, .populate = false }); defer map.destroy(io); var file = try safetensors.File.parse(gpa, map.memory); defer file.deinit();
spec.tensor_prefix = if (file.has("embeddings.word_embeddings.weight")) "" else if (file.has("bert.embeddings.word_embeddings.weight")) "bert." else return error.MissingTensor;
return loadTensors(gpa, &file, spec, cfg, tokenizer); }
fn loadTensors(gpa: Allocator, file: *const safetensors.File, spec: Spec, cfg: Config, tokenizer: Tokenizer) !Model { const h = spec.hidden; const precision: Precision = switch (cfg.precision) { .f32 => .f32, .bf16 => if (bf16.selfTest()) .bf16 else .f32, .bf16_emulated => .bf16_emulated, .int8 => if (int8.selfTest()) .int8 else .f32, .int8_emulated => .int8_emulated, .fastest => if (bf16.selfTest()) .bf16 else if (int8.selfTest()) .int8 else .f32, }; const t: Tensors = .{ .file = file, .prefix = spec.tensor_prefix };
const word = try t.read(gpa, "embeddings.word_embeddings.weight", spec.vocab * h); errdefer gpa.free(word); const pos = try t.read(gpa, "embeddings.position_embeddings.weight", spec.max_positions * h); errdefer gpa.free(pos); const types = try t.read(gpa, "embeddings.token_type_embeddings.weight", spec.type_vocab * h); defer gpa.free(types); for (0..spec.max_positions) |p| for (pos[p * h ..][0..h], types[0..h]) |*a, b| { a.* += b; };
const layers = try gpa.alloc(Layer, spec.layers); errdefer gpa.free(layers); var loaded: usize = 0; errdefer for (layers[0..loaded]) |*layer| freeLayer(gpa, layer); for (layers, 0..) |*layer, l| { layer.* = try loadLayer(gpa, precision, t, l, spec); loaded += 1; }
const emb_ln_g = try t.read(gpa, "embeddings.LayerNorm.weight", h); errdefer gpa.free(emb_ln_g); return .{ .spec = spec, .tokenizer = tokenizer, .word = word, .pos = pos, .emb_ln_g = emb_ln_g, .emb_ln_b = try t.read(gpa, "embeddings.LayerNorm.bias", h), .layers = layers, .precision = precision, }; }
fn loadLayer(gpa: Allocator, precision: Precision, t: Tensors, l: usize, spec: Spec) !Layer { const h = spec.hidden; var layer: Layer = undefined; { const qkv_w = try gpa.alloc(f32, 3 * h * h); defer gpa.free(qkv_w); const qkv_b = try gpa.alloc(f32, 3 * h); defer gpa.free(qkv_b); for ([_][]const u8{ "query", "key", "value" }, 0..) |part, i| { try t.readLayer(l, "attention.self", part, "weight", qkv_w[i * h * h ..][0 .. h * h]); try t.readLayer(l, "attention.self", part, "bias", qkv_b[i * h ..][0..h]); } layer.qkv = try packLinear(gpa, precision, qkv_w, qkv_b, h, 3 * h); } errdefer layer.qkv.deinit(gpa); layer.attn_out = try t.linear(gpa, precision, l, "attention.output.dense", h, h); errdefer layer.attn_out.deinit(gpa); layer.ffn_in = try t.linear(gpa, precision, l, "intermediate.dense", h, spec.intermediate); errdefer layer.ffn_in.deinit(gpa); layer.ffn_out = try t.linear(gpa, precision, l, "output.dense", spec.intermediate, h); errdefer layer.ffn_out.deinit(gpa); layer.ln1_g = try t.layerNew(gpa, l, "attention.output.LayerNorm", "weight", h); errdefer gpa.free(layer.ln1_g); layer.ln1_b = try t.layerNew(gpa, l, "attention.output.LayerNorm", "bias", h); errdefer gpa.free(layer.ln1_b); layer.ln2_g = try t.layerNew(gpa, l, "output.LayerNorm", "weight", h); errdefer gpa.free(layer.ln2_g); layer.ln2_b = try t.layerNew(gpa, l, "output.LayerNorm", "bias", h); return layer; }
pub fn deinit(self: *Model, gpa: Allocator) void { for (self.layers) |*layer| freeLayer(gpa, layer); gpa.free(self.layers); gpa.free(self.word); gpa.free(self.pos); gpa.free(self.emb_ln_g); gpa.free(self.emb_ln_b); self.tokenizer.deinit(gpa); }
pub fn dim(self: *const Model) usize { return self.spec.hidden; }
/// Embed each text into a row of `out` (texts.len × dim floats), pooled and /// (if the model says so) L2-normalized the way sentence-transformers does. pub fn embed(self: *const Model, gpa: Allocator, texts: []const []const u8, out: []f32) !void { const s = self.spec; const h = s.hidden; std.debug.assert(out.len == texts.len * h); if (texts.len == 0) return;
var arena_state = std.heap.ArenaAllocator.init(gpa); defer arena_state.deinit(); const arena = arena_state.allocator();
const prof = self.profile; var t: u64 = if (prof) |pr| pr.now() else 0;
const lens = try arena.alloc(usize, texts.len); var ids: std.ArrayList(u32) = .empty; var max_len: usize = 0; var scratch: std.ArrayList(u8) = .empty; const seq_ids = try arena.alloc(u32, s.max_seq); for (texts, lens) |text, *len| { try scratch.resize(arena, 4 * text.len + 4); len.* = self.tokenizer.encode(text, scratch.items, seq_ids); try ids.appendSlice(arena, seq_ids[0..len.*]); max_len = @max(max_len, len.*); } const rows = ids.items.len; if (prof) |pr| { t = pr.add(.tokenize, t); _ = pr.tokens.fetchAdd(rows, .monotonic); }
const x = try arena.alloc(f32, rows * h); const tmp = try arena.alloc(f32, rows * h); const qkv = try arena.alloc(f32, rows * 3 * h); const ffn = try arena.alloc(f32, rows * s.intermediate); const scores = try arena.alloc(f32, ops.attentionScratch(max_len));
var r: usize = 0; for (lens) |len| { for (0..len) |p| { const row = x[(r + p) * h ..][0..h]; const tok = ids.items[r + p]; for (row, self.word[tok * h ..][0..h], self.pos[p * h ..][0..h]) |*d, w, ps| d.* = w + ps; ops.layerNormRow(row, self.emb_ln_g, self.emb_ln_b, s.eps); } r += len; } if (prof) |pr| t = pr.add(.embeddings, t);
const n = texts.len; const pooled = try arena.alloc(f32, n * h); switch (s.pooling) { .mean => { for (self.layers) |*layer| t = self.encoderLayer(layer, x, tmp, qkv, ffn, scores, lens, rows, t); r = 0; for (lens, 0..) |len, i| { const dst = pooled[i * h ..][0..h]; @memset(dst, 0); for (0..len) |p| for (dst, x[(r + p) * h ..][0..h]) |*d, v| { d.* += v; }; const inv = 1.0 / @as(f32, @floatFromInt(len)); for (dst) |*d| d.* *= inv; r += len; } }, .cls => { for (self.layers[0 .. self.layers.len - 1]) |*layer| t = self.encoderLayer(layer, x, tmp, qkv, ffn, scores, lens, rows, t); // CLS pooling reads only row 0 of the last layer, so after its K/V // (which every query attends over) the rest runs on one row per text const last = &self.layers[self.layers.len - 1]; last.qkv.forward(x, qkv, rows); if (prof) |pr| t = pr.add(.qkv, t); const cls_ctx = try arena.alloc(f32, n * h); r = 0; for (lens, 0..) |len, i| { @memcpy(pooled[i * h ..][0..h], x[r * h ..][0..h]); ops.attention(qkv[r * 3 * h ..], cls_ctx[i * h ..], scores, len, h, s.heads, 1); r += len; } if (prof) |pr| t = pr.add(.attention, t); t = self.feedForward(last, pooled, cls_ctx, tmp, ffn, n, t); }, }
for (0..n) |i| { const row = pooled[i * h ..][0..h]; var scale: f32 = 1; if (s.normalize) { var norm: f32 = 0; for (row) |v| norm += v * v; scale = 1.0 / @max(@sqrt(norm), 1e-12); } for (out[i * h ..][0..h], row) |*o, v| o.* = v * scale; } if (prof) |pr| _ = pr.add(.pool, t); }
/// One full encoder layer over all rows, in place on `x`. fn encoderLayer(self: *const Model, l: *const Layer, x: []f32, tmp: []f32, qkv: []f32, ffn: []f32, scores: []f32, lens: []const usize, rows: usize, t_in: u64) u64 { const h = self.spec.hidden; const prof = self.profile; var t = t_in; l.qkv.forward(x, qkv, rows); if (prof) |pr| t = pr.add(.qkv, t); var r: usize = 0; for (lens) |len| { ops.attention(qkv[r * 3 * h ..], tmp[r * h ..], scores, len, h, self.spec.heads, len); r += len; } if (prof) |pr| t = pr.add(.attention, t); return self.feedForward(l, x, tmp, qkv, ffn, rows, t); }
/// Attention output projection, residual LayerNorm, FFN, residual /// LayerNorm, for `rows` rows of `x` given their attention context `ctx`. /// `tmp` needs rows × hidden floats. fn feedForward(self: *const Model, l: *const Layer, x: []f32, ctx: []const f32, tmp: []f32, ffn: []f32, rows: usize, t_in: u64) u64 { const s = self.spec; const prof = self.profile; var t = t_in; l.attn_out.forward(ctx, tmp, rows); if (prof) |pr| t = pr.add(.attn_out, t); ops.addLayerNorm(x, tmp, l.ln1_g, l.ln1_b, rows, s.eps); if (prof) |pr| t = pr.add(.layernorm, t); l.ffn_in.forward(x, ffn, rows); if (prof) |pr| t = pr.add(.ffn_in, t); ops.gelu(ffn[0 .. rows * s.intermediate]); if (prof) |pr| t = pr.add(.gelu, t); l.ffn_out.forward(ffn, tmp, rows); if (prof) |pr| t = pr.add(.ffn_out, t); ops.addLayerNorm(x, tmp, l.ln2_g, l.ln2_b, rows, s.eps); if (prof) |pr| t = pr.add(.layernorm, t); return t; }};
fn freeLayer(gpa: Allocator, layer: *Layer) void { layer.qkv.deinit(gpa); layer.attn_out.deinit(gpa); layer.ffn_in.deinit(gpa); layer.ffn_out.deinit(gpa); for ([_][]f32{ layer.ln1_g, layer.ln1_b, layer.ln2_g, layer.ln2_b }) |s| gpa.free(s);}
/// Tensor lookup that applies the checkpoint's name prefix.const Tensors = struct { file: *const safetensors.File, prefix: []const u8,
fn read(self: Tensors, gpa: Allocator, name: []const u8, len: usize) ![]f32 { var buf: [160]u8 = undefined; const full = try std.fmt.bufPrint(&buf, "{s}{s}", .{ self.prefix, name }); const out = try gpa.alloc(f32, len); errdefer gpa.free(out); try self.file.readF32(full, out); return out; }
fn readLayer(self: Tensors, layer: usize, block: []const u8, part: []const u8, kind: []const u8, out: []f32) !void { var buf: [160]u8 = undefined; const full = try std.fmt.bufPrint(&buf, "{s}encoder.layer.{d}.{s}.{s}.{s}", .{ self.prefix, layer, block, part, kind }); try self.file.readF32(full, out); }
fn layerNew(self: Tensors, gpa: Allocator, layer: usize, block: []const u8, kind: []const u8, len: usize) ![]f32 { var buf: [160]u8 = undefined; const full = try std.fmt.bufPrint(&buf, "{s}encoder.layer.{d}.{s}.{s}", .{ self.prefix, layer, block, kind }); const out = try gpa.alloc(f32, len); errdefer gpa.free(out); try self.file.readF32(full, out); return out; }
fn linear(self: Tensors, gpa: Allocator, precision: Precision, layer: usize, block: []const u8, in: usize, out: usize) !ops.Linear { const w = try self.layerNew(gpa, layer, block, "weight", in * out); defer gpa.free(w); const b = try self.layerNew(gpa, layer, block, "bias", out); defer gpa.free(b); return packLinear(gpa, precision, w, b, in, out); }};
fn packLinear(gpa: Allocator, precision: Precision, w: []const f32, b: []const f32, in: usize, out: usize) !ops.Linear { return switch (precision) { .f32 => ops.Linear.pack(gpa, w, b, in, out), .bf16 => ops.Linear.packBf16(gpa, w, b, in, out, true), .bf16_emulated => ops.Linear.packBf16(gpa, w, b, in, out, false), .int8 => ops.Linear.packInt8(gpa, w, b, in, out, true), .int8_emulated => ops.Linear.packInt8(gpa, w, b, in, out, false), .fastest => unreachable, };}
// A tiny synthetic snapshot: 32-wide, one head, two layers, six-token vocab.const TestSnapshot = struct { const hidden = 32; const layers = 2; const vocab = "[PAD]\n[UNK]\n[CLS]\n[SEP]\nhello\nworld\n"; const vocab_len = 6; const positions = 8;
fn write(dir: Io.Dir, pooling: []const u8, omit: ?[]const u8) !void { const io = std.testing.io; const a = std.testing.allocator; try dir.writeFile(io, .{ .sub_path = "config.json", .data = std.fmt.comptimePrint( "{{\"architectures\":[\"BertModel\"],\"hidden_act\":\"gelu\",\"hidden_size\":{d},\"num_attention_heads\":1,\"num_hidden_layers\":{d},\"intermediate_size\":{d},\"max_position_embeddings\":{d},\"vocab_size\":{d},\"type_vocab_size\":2}}", .{ hidden, layers, hidden, positions, vocab_len }, ) }); try dir.createDirPath(io, "1_Pooling"); const pool = try std.fmt.allocPrint(a, "{{\"pooling_mode_{s}\":true}}", .{pooling}); defer a.free(pool); try dir.writeFile(io, .{ .sub_path = "1_Pooling/config.json", .data = pool }); try dir.writeFile(io, .{ .sub_path = "modules.json", .data = "[{\"type\":\"sentence_transformers.models.Normalize\"}]" }); try dir.writeFile(io, .{ .sub_path = "tokenizer.json", .data = spec_mod.test_tokenizer }); try dir.writeFile(io, .{ .sub_path = "vocab.txt", .data = vocab });
var names: std.ArrayList([]const u8) = .empty; defer { for (names.items) |n| a.free(n); names.deinit(a); } var sizes: std.ArrayList(usize) = .empty; defer sizes.deinit(a); const add = struct { fn f(al: Allocator, ns: *std.ArrayList([]const u8), ss: *std.ArrayList(usize), name: []const u8, size: usize) !void { try ns.append(al, try al.dupe(u8, name)); try ss.append(al, size); } }.f; try add(a, &names, &sizes, "embeddings.word_embeddings.weight", vocab_len * hidden); try add(a, &names, &sizes, "embeddings.position_embeddings.weight", positions * hidden); try add(a, &names, &sizes, "embeddings.token_type_embeddings.weight", 2 * hidden); try add(a, &names, &sizes, "embeddings.LayerNorm.weight", hidden); try add(a, &names, &sizes, "embeddings.LayerNorm.bias", hidden); for (0..layers) |l| { const blocks = [_][]const u8{ "attention.self.query", "attention.self.key", "attention.self.value", "attention.output.dense", "intermediate.dense", "output.dense", "attention.output.LayerNorm", "output.LayerNorm" }; for (blocks, 0..) |b, i| { const w = try std.fmt.allocPrint(a, "encoder.layer.{d}.{s}.weight", .{ l, b }); defer a.free(w); try add(a, &names, &sizes, w, if (i < 6) hidden * hidden else hidden); const bn = try std.fmt.allocPrint(a, "encoder.layer.{d}.{s}.bias", .{ l, b }); defer a.free(bn); try add(a, &names, &sizes, bn, hidden); } }
var header: std.ArrayList(u8) = .empty; defer header.deinit(a); var data: std.ArrayList(u8) = .empty; defer data.deinit(a); try header.append(a, '{'); var first = true; for (names.items, sizes.items) |name, size| { if (omit) |o| if (std.mem.eql(u8, o, name)) continue; const start = data.items.len; const is_ln_weight = std.mem.endsWith(u8, name, "LayerNorm.weight"); for (0..size) |i| { const v: f32 = if (is_ln_weight) 1 else @sin(@as(f32, @floatFromInt(i + start)) * 0.37) * 0.2; try data.appendSlice(a, std.mem.asBytes(&v)); } if (!first) try header.append(a, ','); first = false; try header.print(a, "\"{s}\":{{\"dtype\":\"F32\",\"shape\":[{d}],\"data_offsets\":[{d},{d}]}}", .{ name, size, start, data.items.len }); } try header.append(a, '}'); var file: std.ArrayList(u8) = .empty; defer file.deinit(a); var len: [8]u8 = undefined; std.mem.writeInt(u64, &len, header.items.len, .little); try file.appendSlice(a, &len); try file.appendSlice(a, header.items); try file.appendSlice(a, data.items); try dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = file.items }); }};
test "a snapshot's pooling mode is honored and outputs are unit length" { var cls_out: [2 * TestSnapshot.hidden]f32 = undefined; var mean_out: [2 * TestSnapshot.hidden]f32 = undefined; for ([_][]const u8{ "cls_token", "mean_tokens" }, [_]*[2 * TestSnapshot.hidden]f32{ &cls_out, &mean_out }) |mode, out| { var tmp = std.testing.tmpDir(.{}); defer tmp.cleanup(); try TestSnapshot.write(tmp.dir, mode, null); var model = try Model.loadDir(std.testing.allocator, std.testing.io, tmp.dir, .{}); defer model.deinit(std.testing.allocator); try std.testing.expectEqual(@as(usize, TestSnapshot.layers), model.layers.len); try model.embed(std.testing.allocator, &.{ "hello world", "world" }, out); for (0..2) |i| { var norm: f32 = 0; for (out[i * TestSnapshot.hidden ..][0..TestSnapshot.hidden]) |v| norm += v * v; try std.testing.expectApproxEqAbs(@as(f32, 1), norm, 1e-5); } } try std.testing.expect(!std.mem.eql(f32, &cls_out, &mean_out));}
test "a snapshot missing a tensor fails to load without leaking" { var tmp = std.testing.tmpDir(.{}); defer tmp.cleanup(); try TestSnapshot.write(tmp.dir, "cls_token", "encoder.layer.1.output.LayerNorm.bias"); try std.testing.expectError(error.MissingTensor, Model.loadDir(std.testing.allocator, std.testing.io, tmp.dir, .{}));}