pyrox.dev
1{
2 lib,
3 stdenv,
4 fetchFromGitHub,
5 cmake,
6 llvmPackages, # openmp
7 withMkl ? false,
8 mkl,
9 withCUDA ? false,
10 withCuDNN ? false,
11 cudaPackages,
12 # Enabling both withOneDNN and withOpenblas is broken
13 # https://github.com/OpenNMT/CTranslate2/issues/1294
14 withOneDNN ? false,
15 oneDNN,
16 withOpenblas ? true,
17 openblas,
18 withRuy ? true,
19
20 # passthru tests
21 libretranslate,
22 wyoming-faster-whisper,
23}:
24
25let
26 cmakeBool = b: if b then "ON" else "OFF";
27 stdenv' = if withCUDA then cudaPackages.backendStdenv else stdenv;
28in
29stdenv'.mkDerivation rec {
30 pname = "ctranslate2";
31 version = "4.6.0";
32
33 src = fetchFromGitHub {
34 owner = "OpenNMT";
35 repo = "CTranslate2";
36 rev = "v${version}";
37 hash = "sha256-EM2kunqtxo0BTIzrEomfaRsdav7sx6QEOhjDtjjSoYY=";
38 fetchSubmodules = true;
39 };
40
41 nativeBuildInputs = [
42 cmake
43 ]
44 ++ lib.optionals withCUDA [
45 cudaPackages.cuda_nvcc
46 ];
47
48 cmakeFlags = [
49 # https://opennmt.net/CTranslate2/installation.html#build-options
50 # https://github.com/OpenNMT/CTranslate2/blob/54810350e662ebdb01ecbf8e4a746f02aeff1dd7/python/tools/prepare_build_environment_linux.sh#L53
51 # https://github.com/OpenNMT/CTranslate2/blob/59d223abcc7e636c1c2956e62482bc3299cc7766/python/tools/prepare_build_environment_macos.sh#L12
52 "-DOPENMP_RUNTIME=COMP"
53 "-DWITH_CUDA=${cmakeBool withCUDA}"
54 "-DWITH_CUDNN=${cmakeBool withCuDNN}"
55 "-DWITH_DNNL=${cmakeBool withOneDNN}"
56 "-DWITH_OPENBLAS=${cmakeBool withOpenblas}"
57 "-DWITH_RUY=${cmakeBool withRuy}"
58 "-DWITH_MKL=${cmakeBool withMkl}"
59 ]
60 ++ lib.optional stdenv.hostPlatform.isDarwin "-DWITH_ACCELERATE=ON";
61
62 buildInputs =
63 lib.optionals withMkl [
64 mkl
65 ]
66 ++ lib.optionals withCUDA [
67 cudaPackages.cuda_cccl # <nv/target> required by the fp16 headers in cudart
68 cudaPackages.cuda_cudart
69 cudaPackages.libcublas
70 cudaPackages.libcurand
71 ]
72 ++ lib.optionals (withCUDA && withCuDNN) [
73 cudaPackages.cudnn
74 ]
75 ++ lib.optionals withOneDNN [
76 oneDNN
77 ]
78 ++ lib.optionals withOpenblas [
79 openblas
80 ]
81 ++ lib.optionals stdenv.hostPlatform.isDarwin [
82 llvmPackages.openmp
83 ];
84
85 passthru.tests = {
86 inherit
87 libretranslate
88 wyoming-faster-whisper
89 ;
90 };
91
92 meta = with lib; {
93 description = "Fast inference engine for Transformer models";
94 mainProgram = "ct2-translator";
95 homepage = "https://github.com/OpenNMT/CTranslate2";
96 changelog = "https://github.com/OpenNMT/CTranslate2/blob/${src.rev}/CHANGELOG.md";
97 license = licenses.mit;
98 maintainers = with maintainers; [
99 hexa
100 misuzu
101 ];
102 broken = !(withCuDNN -> withCUDA);
103 };
104}