tjh.dev
1{ lib
2, fetchFromGitHub
3, buildPythonPackage
4, substituteAll
5, cudaSupport ? false
6
7# runtime
8, ffmpeg
9
10# propagates
11, numpy
12, torch
13, torchWithCuda
14, tqdm
15, more-itertools
16, transformers
17, ffmpeg-python
18, numba
19, openai-triton
20, scipy
21, tiktoken
22
23# tests
24, pytestCheckHook
25}:
26
27buildPythonPackage rec {
28 pname = "whisper";
29 version = "20230314";
30 format = "setuptools";
31
32 src = fetchFromGitHub {
33 owner = "openai";
34 repo = pname;
35 rev = "refs/tags/v${version}";
36 hash = "sha256-qQCELjRFeRCT1k1CBc3netRtFvt+an/EbkrgnmiX/mc=";
37 };
38
39 patches = [
40 (substituteAll {
41 src = ./ffmpeg-path.patch;
42 inherit ffmpeg;
43 })
44 ];
45
46 propagatedBuildInputs = [
47 numpy
48 tqdm
49 more-itertools
50 transformers
51 ffmpeg-python
52 numba
53 scipy
54 tiktoken
55 ] ++ lib.optionals (!cudaSupport) [
56 torch
57 ] ++ lib.optionals (cudaSupport) [
58 openai-triton
59 torchWithCuda
60 ];
61
62 postPatch = ''
63 substituteInPlace requirements.txt \
64 --replace "tiktoken==0.3.1" "tiktoken>=0.3.1"
65 ''
66 # openai-triton is only needed for CUDA support.
67 # triton needs CUDA to be build.
68 # -> by making it optional, we can build whisper without unfree packages enabled
69 + lib.optionalString (!cudaSupport) ''
70 sed -i '/if sys.platform.startswith("linux") and platform.machine() == "x86_64":/{N;d}' setup.py
71 '';
72
73 preCheck = ''
74 export HOME=$TMPDIR
75 '';
76
77 nativeCheckInputs = [
78 pytestCheckHook
79 ];
80
81 disabledTests = [
82 # requires network access to download models
83 "test_tokenizer"
84 "test_transcribe"
85 # requires NVIDIA drivers
86 "test_dtw_cuda_equivalence"
87 "test_median_filter_equivalence"
88 ];
89
90 meta = with lib; {
91 changelog = "https://github.com/openai/whisper/blob/v$[version}/CHANGELOG.md";
92 description = "General-purpose speech recognition model";
93 homepage = "https://github.com/openai/whisper";
94 license = licenses.mit;
95 maintainers = with maintainers; [ hexa MayNiklas ];
96 };
97}