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