Skip to content

Commit

Permalink
build(nix): Exclude gguf-py from devShells
Browse files Browse the repository at this point in the history
  • Loading branch information
ditsuke committed Feb 23, 2024
1 parent 2b51e13 commit 1d0a97c
Showing 1 changed file with 24 additions and 15 deletions.
39 changes: 24 additions & 15 deletions .devops/nix/package.nix
Original file line number Diff line number Diff line change
Expand Up @@ -63,24 +63,27 @@ let
suffices != [ ]
) ", accelerated with ${strings.concatStringsSep ", " suffices}";

mapToPythonPackages = ps: packages: map (package: ps.${package}) packages;

# TODO: package the Python in this repository in a Nix-like way.
# It'd be nice to migrate to buildPythonPackage, as well as ensure this repo
# is PEP 517-compatible, and ensure the correct .dist-info is generated.
# https://peps.python.org/pep-0517/
llama-python = python3.withPackages (ps: [
ps.numpy
ps.sentencepiece
gguf-py
]);
llama-python-base-deps = [
"numpy"
"sentencepiece"
];

# TODO(Green-Sky): find a better way to opt-into the heavy ml python runtime
llama-python-extra = python3.withPackages (ps: [
ps.numpy
ps.sentencepiece
ps.tiktoken
ps.torchWithoutCuda
ps.transformers
]);
llama-python-full-deps = llama-python-base-deps ++ [
"tiktoken"
"torchWithoutCuda"
"transformers"
];

llama-python-base-with-gguf = python3.withPackages (
ps: (mapToPythonPackages ps llama-python-base-deps) ++ [ gguf-py ]
);

# apple_sdk is supposed to choose sane defaults, no need to handle isAarch64
# separately
Expand Down Expand Up @@ -148,7 +151,7 @@ effectiveStdenv.mkDerivation (finalAttrs: {
# TODO: Package up each Python script or service appropriately.
# If we were to migrate to buildPythonPackage and prepare the `pyproject.toml`,
# we could make those *.py into setuptools' entrypoints
substituteInPlace ./*.py --replace "/usr/bin/env python" "${llama-python}/bin/python"
substituteInPlace ./*.py --replace "/usr/bin/env python" "${llama-python-base-with-gguf}/bin/python"
'';

nativeBuildInputs =
Expand Down Expand Up @@ -233,7 +236,10 @@ effectiveStdenv.mkDerivation (finalAttrs: {
shell = mkShell {
name = "shell-${finalAttrs.finalPackage.name}";
description = "contains numpy and sentencepiece";
buildInputs = [ llama-python ];
buildInputs = [
python3.withPackages
(ps: mapToPythonPackages ps llama-python-base-deps)
];
inputsFrom = [ finalAttrs.finalPackage ];
shellHook = ''
addToSearchPath "LD_LIBRARY_PATH" "${lib.getLib effectiveStdenv.cc.cc}/lib"
Expand All @@ -243,7 +249,10 @@ effectiveStdenv.mkDerivation (finalAttrs: {
shell-extra = mkShell {
name = "shell-extra-${finalAttrs.finalPackage.name}";
description = "contains numpy, sentencepiece, torchWithoutCuda, and transformers";
buildInputs = [ llama-python-extra ];
buildInputs = [
python3.withPackages
(ps: mapToPythonPackages ps llama-python-full-deps)
];
inputsFrom = [ finalAttrs.finalPackage ];
};
};
Expand Down

0 comments on commit 1d0a97c

Please sign in to comment.