Read grammar vocabulary from llm_config for Nemotron Omni

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));
+    }
+}

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");

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;