diff --git a/zig/src/server/grammar.zig b/zig/src/server/grammar.zig index 543f08a265..8799e14aa3 100644 --- a/zig/src/server/grammar.zig +++ b/zig/src/server/grammar.zig @@ -13,6 +13,16 @@ pub const fields = [_][]const u8{ "response_format", "guided_json", "guided_rege pub const Kind = @FieldType(api.Structure, "kind"); pub const Spec = struct { kind: Kind, text: []const u8 = "", field: []const u8 = "response_format" }; +/// The text vocabulary within the grammar ABI range: text_config, llm_config, then root. +pub fn vocabSize(cfg: Value) ?u32 { + for ([_]?Value{ cfg.get("text_config"), cfg.get("llm_config"), cfg }) |part| if (part) |p| { + const v = p.get("vocab_size") orelse continue; + const n = v.int64() orelse continue; + if (n > 0 and n <= std.math.maxInt(c_int)) return @intCast(n); + }; + return null; +} + /// The JSON decoder's own message, without its position. fn bareMessage(full: []const u8) []const u8 { const at = std.mem.lastIndexOf(u8, full, ": line ") orelse return full; @@ -99,3 +109,43 @@ pub fn refusal(cx: *Cx, body: Value) errors.Refused!void { if (tools != null and tools.?.truthy() and try tool_specs.choiceRequiresCall(cx.a, body.get("tool_choice"))) return cx.fail(.request, "{s} cannot be combined with tool_choice \"required\" or a named function: send one", .{s.field}); } + +test "grammar vocabulary reads Nemotron Omni's text model" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const cfg = (try json.parse(a, + \\{"model_type":"nemotron_h_omni","llm_config":{"model_type":"nemotron_h","vocab_size":131072}} + )).ok; + try std.testing.expectEqual(@as(?u32, 131072), vocabSize(cfg)); +} + +test "grammar vocabulary preserves nested text and root precedence" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const cases = [_]struct { []const u8, ?u32 }{ + .{ "{\"vocab_size\":32000}", 32000 }, + .{ "{\"text_config\":{\"vocab_size\":32001},\"vocab_size\":32000}", 32001 }, + .{ "{\"text_config\":{},\"vocab_size\":32000}", 32000 }, + .{ "{\"text_config\":null,\"vocab_size\":32000}", 32000 }, + .{ "{\"llm_config\":{\"vocab_size\":131072},\"vocab_size\":32000}", 131072 }, + .{ "{\"llm_config\":null,\"vocab_size\":32000}", 32000 }, + .{ "{\"text_config\":{\"vocab_size\":-1},\"llm_config\":{\"vocab_size\":131072}}", 131072 }, + .{ "{\"vocab_size\":2147483647}", 2147483647 }, + .{ "{\"vocab_size\":2147483648}", null }, + .{ "{\"vocab_size\":4294967295}", null }, + .{ "{\"text_config\":{\"vocab_size\":2147483648},\"llm_config\":{\"vocab_size\":131072}}", 131072 }, + .{ "{\"text_config\":{\"vocab_size\":32001},\"llm_config\":{\"vocab_size\":131072},\"vocab_size\":32000}", 32001 }, + .{ "{}", null }, + .{ "{\"vocab_size\":-1}", null }, + .{ "{\"vocab_size\":0}", null }, + .{ "{\"vocab_size\":4294967296}", null }, + .{ "{\"vocab_size\":\"131072\"}", null }, + .{ "{\"vocab_size\":true}", null }, + }; + for (cases) |c| { + const cfg = (try json.parse(a, c[0])).ok; + try std.testing.expectEqual(c[1], vocabSize(cfg)); + } +} diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 381c864ad0..47f3d6bf1e 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -25,6 +25,7 @@ test { _ = @import("tool_stream.zig"); _ = @import("tool_params.zig"); _ = @import("tool_grammar.zig"); + _ = @import("grammar.zig"); _ = @import("pyrepr.zig"); _ = @import("metrics.zig"); _ = @import("http_conn.zig"); diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 5406581774..fc6c1a8da2 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -229,7 +229,7 @@ pub const Server = struct { return .{ .kind = kind, .text = text, .after = if (thinking) srv.text.tokenId(srv.markers.close) else null, .compiler = c }; } - /// The model's grammar compiler: tokenizer.json, config.json's vocab_size (text_config first), the eos ids. + /// The model's grammar compiler: tokenizer.json, the text model's vocab_size, the eos ids. fn grammarCompiler(srv: *Server) ?*api.grammar.Compiler { srv.grammar_mutex.lockUncancelable(srv.io); defer srv.grammar_mutex.unlock(srv.io); @@ -244,19 +244,15 @@ pub const Server = struct { .ok => |v| v, .err => break :blk "config.json is not JSON", }; - var vocab: ?u32 = null; - for ([_]?json.Value{ cfg.get("text_config"), cfg }) |part| if (vocab == null) if (part) |p| if (p.get("vocab_size")) |v| if (v.int64()) |n| { - vocab = @intCast(n); - }; - if (vocab == null) break :blk "structured output needs the checkpoint's config.json vocab_size"; + const vocab = grammar.vocabSize(cfg) orelse break :blk "structured output needs the checkpoint's config.json vocab_size"; var err: [512]u8 = undefined; const c = a.create(api.grammar.Compiler) catch break :blk "out of memory"; - c.* = api.grammar.Compiler.open(srv.io, a, tok, vocab.?, srv.eos, &err) catch |e| break :blk switch (e) { + c.* = api.grammar.Compiler.open(srv.io, a, tok, vocab, srv.eos, &err) catch |e| break :blk switch (e) { error.NoGrammarLibrary => "the server cannot load " ++ api.grammar.library ++ " (xgrammar; zig build grammar)", else => a.dupe(u8, std.mem.sliceTo(&err, 0)) catch "out of memory", }; srv.grammars = c; - log.line("structured output: xgrammar over {d} tokens", .{vocab.?}); + log.line("structured output: xgrammar over {d} tokens", .{vocab}); break :blk null; }; return srv.grammars;