tjh.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 pydantic,
15 sentencepiece,
16 tiktoken,
17 torch,
18 transformers,
19 triton,
20
21 # tests
22 pytestCheckHook,
23 writableTmpDirAsHomeHook,
24}:
25
26buildPythonPackage rec {
27 pname = "xgrammar";
28 version = "0.1.22";
29 pyproject = true;
30
31 src = fetchFromGitHub {
32 owner = "mlc-ai";
33 repo = "xgrammar";
34 tag = "v${version}";
35 fetchSubmodules = true;
36 hash = "sha256-mz6eabETkAzDoPjXE5VJvgrR1vnXXmx3JO4xZRH4TRQ=";
37 };
38
39 patches = [
40 ./0001-fix-find-nanobind-from-python-module.patch
41 ];
42
43 build-system = [
44 cmake
45 ninja
46 nanobind
47 scikit-build-core
48 ];
49 dontUseCmakeConfigure = true;
50
51 dependencies = [
52 pydantic
53 sentencepiece
54 tiktoken
55 torch
56 transformers
57 ]
58 ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isx86_64) [
59 triton
60 ];
61
62 nativeCheckInputs = [
63 pytestCheckHook
64 writableTmpDirAsHomeHook
65 ];
66
67 NIX_CFLAGS_COMPILE = toString [
68 # xgrammar hardcodes -flto=auto while using static linking, which can cause linker errors without this additional flag.
69 "-ffat-lto-objects"
70 ];
71
72 disabledTests = [
73 # You are trying to access a gated repo.
74 "test_grammar_compiler"
75 "test_grammar_matcher"
76 "test_grammar_matcher_ebnf"
77 "test_grammar_matcher_json"
78 "test_grammar_matcher_json_schema"
79 "test_grammar_matcher_tag_dispatch"
80 "test_regex_converter"
81 "test_tokenizer_info"
82
83 # Torch not compiled with CUDA enabled
84 "test_token_bitmask_operations"
85
86 # AssertionError
87 "test_json_schema_converter"
88 ];
89
90 pythonImportsCheck = [ "xgrammar" ];
91
92 meta = {
93 description = "Efficient, Flexible and Portable Structured Generation";
94 homepage = "https://xgrammar.mlc.ai";
95 changelog = "https://github.com/mlc-ai/xgrammar/releases/tag/${src.tag}";
96 license = lib.licenses.asl20;
97 };
98}