search

mlc-ai / xgrammar

Efficient, Flexible and Portable Structured Generation
https://xgrammar.mlc.ai/
Apache License 2.0
380 stars 18 forks source link

[Contrib] Implement HF LogitsProcessor for integration with transformers #68

Closed CharlieFRuan closed 1 week ago

CharlieFRuan commented 1 week ago

This PR adds a contrib folder under python/xgrammar and implements xgrammar.contrib.hf.LogitsProcessor.

We also add examples/hf_transformers on how to use XGrammar in Huggingface's transformers library with a minimal LogitsProcessor. The API is the following

xgr_logits_processor = xgr.contrib.hf.LogitsProcessor(compiled_grammar)
generated_ids = model.generate(
    **model_inputs, logits_processor=[xgr_logits_processor]
)
tqchen commented 1 week ago

Please update to latest api and help testing out

CharlieFRuan commented 1 week ago

Rebased to use new API, tested with single-request and batched inference. Will add more changes in the CPU bitmasking PR.