diff --git a/pkgs/development/python-modules/ctranslate2/default.nix b/pkgs/development/python-modules/ctranslate2/default.nix new file mode 100644 index 000000000000..4cbf0ecd7124 --- /dev/null +++ b/pkgs/development/python-modules/ctranslate2/default.nix @@ -0,0 +1,84 @@ +{ lib +, buildPythonPackage + +# build-system +, pybind11 +, setuptools + +# dependencies +, ctranslate2-cpp +, numpy +, pyyaml + +# tests +, pytestCheckHook +, tensorflow +, torch +, transformers +, wurlitzer +}: + +buildPythonPackage rec { + inherit (ctranslate2-cpp) pname version src; + format = "setuptools"; + + # https://github.com/OpenNMT/CTranslate2/tree/master/python + sourceRoot = "source/python"; + + nativeBuildInputs = [ + pybind11 + setuptools + ]; + + buildInputs = [ + ctranslate2-cpp + ]; + + propagatedBuildInputs = [ + numpy + pyyaml + ]; + + pythonImportsCheck = [ + # https://opennmt.net/CTranslate2/python/overview.html + "ctranslate2" + "ctranslate2.converters" + "ctranslate2.models" + "ctranslate2.specs" + ]; + + nativeCheckInputs = [ + pytestCheckHook + tensorflow + torch + transformers + wurlitzer + ]; + + preCheck = '' + # run tests against build result, not sources + rm -rf ctranslate2 + + export HOME=$TMPDIR + ''; + + disabledTests = [ + # AssertionError: assert 'int8' in {'float32'} + "test_get_supported_compute_types" + ]; + + disabledTestPaths = [ + # TODO: ModuleNotFoundError: No module named 'opennmt' + "tests/test_opennmt_tf.py" + # OSError: We couldn't connect to 'https://huggingface.co' to load this file + "tests/test_transformers.py" + ]; + + meta = with lib; { + description = "Fast inference engine for Transformer models"; + homepage = "https://github.com/OpenNMT/CTranslate2"; + changelog = "https://github.com/OpenNMT/CTranslate2/blob/${src.rev}/CHANGELOG.md"; + license = licenses.mit; + maintainers = with maintainers; [ hexa ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index b9051f34031d..36a081dd6f0a 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -2218,6 +2218,10 @@ self: super: with self; { ctap-keyring-device = callPackage ../development/python-modules/ctap-keyring-device { }; + ctranslate2 = callPackage ../development/python-modules/ctranslate2 { + ctranslate2-cpp = pkgs.ctranslate2; + }; + cu2qu = callPackage ../development/python-modules/cu2qu { }; cucumber-tag-expressions = callPackage ../development/python-modules/cucumber-tag-expressions { };