diff --git a/pkgs/development/python-modules/mlx-lm/default.nix b/pkgs/development/python-modules/mlx-lm/default.nix new file mode 100644 index 0000000000000..e724bb22ff166 --- /dev/null +++ b/pkgs/development/python-modules/mlx-lm/default.nix @@ -0,0 +1,77 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + setuptools, + jinja2, + mlx, + numpy, + protobuf, + pyyaml, + transformers, + sentencepiece, + pytestCheckHook, + writableTmpDirAsHomeHook, +}: + +buildPythonPackage rec { + pname = "mlx-lm"; + version = "0.26.0"; + pyproject = true; + + src = fetchFromGitHub { + owner = "ml-explore"; + repo = "mlx-lm"; + tag = "v${version}"; + hash = "sha256-J69XIqsjQ4sQqhx+EkjKcVXVlQ4A4PGJvICSiCfoSOA="; + }; + + build-system = [ + setuptools + ]; + + dependencies = [ + jinja2 + mlx + numpy + protobuf + pyyaml + transformers + ]; + + nativeCheckInputs = [ + writableTmpDirAsHomeHook + pytestCheckHook + sentencepiece + ]; + + pythonImportsCheck = [ + "mlx_lm" + ]; + + disabledTestPaths = [ + # Requires network access to huggingface.co + "tests/test_datsets.py" + "tests/test_generate.py" + "tests/test_server.py" + "tests/test_tokenizers.py" + "tests/test_utils.py::TestUtils::test_convert" + "tests/test_utils.py::TestUtils::test_load" + "tests/test_utils_load_model.py" + "tests/test_prompt_cache.py::TestPromptCache::test_cache_to_quantized" + "tests/test_prompt_cache.py::TestPromptCache::test_cache_with_generate" + "tests/test_prompt_cache.py::TestPromptCache::test_trim_cache_with_generate" + # RuntimeError: [metal_kernel] No GPU back-end. + "tests/test_models.py::TestModels::test_bitnet" + ]; + + meta = { + description = "Run LLMs with MLX"; + homepage = "https://github.com/ml-explore/mlx-lm"; + license = lib.licenses.mit; + platforms = [ + "aarch64-darwin" + ]; + maintainers = with lib.maintainers; [ ferrine ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 3bd972a556b97..d177c4e62cddc 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -9314,6 +9314,8 @@ self: super: with self; { mlx = callPackage ../development/python-modules/mlx { }; + mlx-lm = callPackage ../development/python-modules/mlx-lm { }; + mlxtend = callPackage ../development/python-modules/mlxtend { }; mmcif-pdbx = callPackage ../development/python-modules/mmcif-pdbx { };