pkgs/development/python-modules/ctranslate2/default.nix

   1 { lib
   2 , buildPythonPackage
   3
   4 # build-system
   5 , pybind11
   6 , setuptools
   7
   8 # dependencies
   9 , ctranslate2-cpp
  10 , numpy
  11 , pyyaml
  12
  13 # tests
  14 , pytestCheckHook
  15 , tensorflow-bin
  16 , torch
  17 , transformers
  18 , wurlitzer
  19 }:
  20
  21 buildPythonPackage rec {
  22   inherit (ctranslate2-cpp) pname version src;
  23   format = "setuptools";
  24
  25   # https://github.com/OpenNMT/CTranslate2/tree/master/python
  26   sourceRoot = "${src.name}/python";
  27
  28   nativeBuildInputs = [
  29     pybind11
  30     setuptools
  31   ];
  32
  33   buildInputs = [
  34     ctranslate2-cpp
  35   ];
  36
  37   propagatedBuildInputs = [
  38     numpy
  39     pyyaml
  40   ];
  41
  42   pythonImportsCheck = [
  43     # https://opennmt.net/CTranslate2/python/overview.html
  44     "ctranslate2"
  45     "ctranslate2.converters"
  46     "ctranslate2.models"
  47     "ctranslate2.specs"
  48   ];
  49
  50   nativeCheckInputs = [
  51     pytestCheckHook
  52     tensorflow-bin
  53     torch
  54     transformers
  55     wurlitzer
  56   ];
  57
  58   preCheck = ''
  59     # run tests against build result, not sources
  60     rm -rf ctranslate2
  61
  62     export HOME=$TMPDIR
  63   '';
  64
  65   disabledTests = [
  66     # AssertionError: assert 'int8' in {'float32'}
  67     "test_get_supported_compute_types"
  68   ];
  69
  70   disabledTestPaths = [
  71     # TODO: ModuleNotFoundError: No module named 'opennmt'
  72     "tests/test_opennmt_tf.py"
  73     # OSError: We couldn't connect to 'https://huggingface.co' to load this file
  74     "tests/test_transformers.py"
  75   ];
  76
  77   meta = with lib; {
  78     description = "Fast inference engine for Transformer models";
  79     homepage = "https://github.com/OpenNMT/CTranslate2";
  80     changelog = "https://github.com/OpenNMT/CTranslate2/blob/${src.rev}/CHANGELOG.md";
  81     license = licenses.mit;
  82     maintainers = with maintainers; [ hexa ];
  83   };
  84 }