Add FunctionCallingLogitsProcessor for tool call support#1910
Add FunctionCallingLogitsProcessor for tool call support#1910Youssefares wants to merge 6 commits into
Conversation
Adds a logits processor that switches between free generation and constrained JSON generation around tool call delimiters (<tool_call> and </tool_call>). Supports batched generation with per-sequence phase tracking and multiple sequential tool calls. fix dottxt-ai#1626
Pulls the output-type-to-logits-processor dispatch out of SteerableGenerator.__init__ into a get_logits_processor helper in the backends module, next to the type-specific get_*_logits_processor functions. FunctionCallingLogitsProcessor now reuses this helper instead of hand-rolling JSON schema conversion, so it inherits support for dataclasses, typed dicts, and nested types. Also renames the tool schema parameter to tools_type.
Mirrors the OutlinesCoreLogitsProcessor dispatch: pick the library-specific masking function once at construction based on tensor_library_name, rather than importing torch/mlx and running isinstance checks on every CLOSING step. The drop mask depends only on the close token and vocab size, so it is built once and cached.
|
📚 Documentation preview: https://dottxt-ai.github.io/outlines/pr-preview/pr-1910/ Preview updates automatically with each commit. |
Annotate _inner_processors as List[Optional[...]] and _close_drop_mask as Optional[np.ndarray], add an assert for type narrowing after the inner processor is created, and mark TensorType indexing/masking ops with type: ignore to match the convention in outlines_core.py.
- Drop the close-mask cache (negligible benefit, removed a branch) - Turn the unreachable `elif CLOSING` into `else` - Mark the MLX masking path `# pragma: no cover` (matches outlines_core) - Add tests for the missing-delimiter error, prototype reuse on non-outlines_core backends, and the no-`_guides` path
|
Hi @Youssefares! Thanks for looking into this topic. I think it's something for which we still need to do a bit more thinking to ensure we use an interface that would work for various tool calling format though. We want to avoid the situation in which we have to break the interface first introduced because it's not exhaustive enough. |
Thanks, Robin. Makes sense! Would it be useful to look into a few formats we're looking to support? and spec'ing out the different possible interfaces. If so, what formats are they? I can look into it and update the issue with details if you tell me 3-4 names that you believe will give us good breadth 🙌 |
|
Thanks for proposing! We ideally want an approach that could accommodate the best open source models such as those from |
|
Np! Appreciate the openness. Sounds good! As a way to leave this in a semi-mergeable state, I will have look at the open source model formats and will update the corresponding issue (if not the PR too). If simple enough I will work into this PR, then we can merge at leisure 🙌 |
2 participants
Summary
Adds a
FunctionCallingLogitsProcessorthat switches between free generation and constrained JSON generation around tool call delimiters, implementing the v0 design from #1626.Refactors get_logits_processor to reuse for compiling tools_type.
How it works
Each sequence in the batch independently transitions through three phases:
OutlinesCoreLogitsProcessorthat constrains generation to the JSON schema. Checksguide.is_finished()after each step.-infexcept the close delimiter.After closing, returns to PASSTHROUGH — supports multiple tool calls per generation.
Implementation details
Indexonce at construction. Index 1:* Batch but Guide 1:1 Batch.input_idsso the inner processor only sees tokens generated since entering ARGUMENTS, not the passthrough prefix or open delimiter._force_single_token_rowhandles numpy, torch, and MLX (MLX arrays are immutable, so standard item assignment doesn't work).Tests
OutlinesCoreLogitsProcessor, GPT-2 tokenizer, and simulated generation loop: verifies the output between delimiters is valid JSON matching the schema.fix #1626