pyrox.dev
1{ lib
2, buildPythonPackage
3, fetchurl
4, protobuf
5, pymorphy3
6, pymorphy3-dicts-uk
7, sentencepiece
8, spacy
9, spacy-pkuseg
10, spacy-transformers
11, writeScript
12, stdenv
13, jq
14, nix
15, moreutils
16}:
17let
18 buildModelPackage = { pname, version, sha256, license }:
19
20 let
21 lang = builtins.substring 0 2 pname;
22 requires-protobuf = pname == "fr_dep_news_trf" || pname == "uk_core_news_trf";
23 in
24 buildPythonPackage {
25 inherit pname version;
26
27 src = fetchurl {
28 url = "https://github.com/explosion/spacy-models/releases/download/${pname}-${version}/${pname}-${version}.tar.gz";
29 inherit sha256;
30 };
31
32 propagatedBuildInputs = [ spacy ]
33 ++ lib.optionals (lib.hasSuffix "_trf" pname) [ spacy-transformers ]
34 ++ lib.optionals (lang == "ru") [ pymorphy3 ]
35 ++ lib.optionals (lang == "uk") [ pymorphy3 pymorphy3-dicts-uk ]
36 ++ lib.optionals (lang == "zh") [ spacy-pkuseg ]
37 ++ lib.optionals (pname == "fr_dep_news_trf") [ sentencepiece ];
38
39 postPatch = lib.optionalString requires-protobuf ''
40 substituteInPlace meta.json \
41 --replace "protobuf<3.21.0" "protobuf"
42 '';
43
44 nativeBuildInputs = lib.optionals requires-protobuf [
45 protobuf
46 ];
47
48 pythonImportsCheck = [ pname ];
49
50 passthru.updateScript = writeScript "update-spacy-models" ''
51 #!${stdenv.shell}
52 set -eou pipefail
53 PATH=${lib.makeBinPath [ jq nix moreutils ]}
54
55 IFS=. read -r major minor patch <<<"${spacy.version}"
56 spacyVersion="$(echo "$major.$minor.0")"
57
58 pushd pkgs/development/python-modules/spacy/ || exit
59
60 jq -r '.[] | .pname' models.json | while IFS= read -r pname; do
61 if [ "$(jq --arg pname "$pname" -r '.[] | select(.pname == $pname) | .version' models.json)" == "$spacyVersion" ]; then
62 continue
63 fi
64
65 newHash="$(nix-prefetch-url "https://github.com/explosion/spacy-models/releases/download/$pname-$spacyVersion/$pname-$spacyVersion.tar.gz")"
66 jq --arg newHash "$newHash" --arg pname "$pname" --arg spacyVersion "$spacyVersion" \
67 '[(.[] | select(.pname != $pname)), (.[] | select(.pname == $pname) | .sha256 = $newHash | .version = $spacyVersion)] | sort_by(.pname)' \
68 models.json | sponge models.json
69 done
70
71 popd || exit
72 '';
73
74 meta = with lib; {
75 description = "Models for the spaCy NLP library";
76 homepage = "https://github.com/explosion/spacy-models";
77 license = licenses.${license};
78 maintainers = with maintainers; [ rvl ];
79 };
80 };
81
82 makeModelSet = models: with lib; listToAttrs (map (m: nameValuePair m.pname (buildModelPackage m)) models);
83
84in
85makeModelSet (lib.importJSON ./models.json)