pyrox.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6
7 # build-system
8 cmake,
9 ninja,
10 nanobind,
11 scikit-build-core,
12
13 # dependencies
14 mlx-lm,
15 pydantic,
16 sentencepiece,
17 tiktoken,
18 torch,
19 transformers,
20 triton,
21
22 # tests
23 pytestCheckHook,
24 writableTmpDirAsHomeHook,
25}:
26
27buildPythonPackage rec {
28 pname = "xgrammar";
29 version = "0.1.24";
30 pyproject = true;
31
32 src = fetchFromGitHub {
33 owner = "mlc-ai";
34 repo = "xgrammar";
35 tag = "v${version}";
36 fetchSubmodules = true;
37 hash = "sha256-K+GCHuWKF449JaGWr7FQrDeJS3pxmVKnGf68L53LrK0=";
38 };
39
40 patches = [
41 ./0001-fix-find-nanobind-from-python-module.patch
42 ];
43
44 build-system = [
45 cmake
46 ninja
47 nanobind
48 scikit-build-core
49 ];
50 dontUseCmakeConfigure = true;
51
52 dependencies = [
53 pydantic
54 sentencepiece
55 tiktoken
56 torch
57 transformers
58 ]
59 ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isx86_64) [
60 triton
61 ]
62 ++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
63 mlx-lm
64 ];
65
66 nativeCheckInputs = [
67 pytestCheckHook
68 writableTmpDirAsHomeHook
69 ];
70
71 NIX_CFLAGS_COMPILE = lib.optionalString stdenv.hostPlatform.isLinux (toString [
72 # xgrammar hardcodes -flto=auto while using static linking, which can cause linker errors without this additional flag.
73 "-ffat-lto-objects"
74 ]);
75
76 disabledTests = [
77 # You are trying to access a gated repo.
78 "test_grammar_compiler"
79 "test_grammar_matcher"
80 "test_grammar_matcher_ebnf"
81 "test_grammar_matcher_json"
82 "test_grammar_matcher_json_schema"
83 "test_grammar_matcher_tag_dispatch"
84 "test_regex_converter"
85 "test_serialize_compiled_grammar_with_hf_tokenizer"
86 "test_tokenizer_info"
87
88 # Torch not compiled with CUDA enabled
89 "test_token_bitmask_operations"
90
91 # AssertionError
92 "test_json_schema_converter"
93 ];
94
95 pythonImportsCheck = [ "xgrammar" ];
96
97 meta = {
98 description = "Efficient, Flexible and Portable Structured Generation";
99 homepage = "https://xgrammar.mlc.ai";
100 changelog = "https://github.com/mlc-ai/xgrammar/releases/tag/${src.tag}";
101 license = lib.licenses.asl20;
102 badPlatforms = [
103 # error: ‘operator delete’ called on unallocated object ‘result’ [-Werror=free-nonheap-object]
104 "aarch64-linux"
105 ];
106 };
107}