pkgs/development/python-modules/spacy/models.nix

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