75 lines
1.5 KiB
Nix
75 lines
1.5 KiB
Nix
{
|
|
pkgs,
|
|
lib,
|
|
python ? pkgs.python312,
|
|
...
|
|
}:
|
|
|
|
let
|
|
pythonWithAI = python.withPackages (
|
|
ps:
|
|
let
|
|
basePackages = with ps; [
|
|
accelerate
|
|
albumentations
|
|
altair
|
|
datasets
|
|
evaluate
|
|
huggingface-hub
|
|
ipykernel
|
|
jupyterlab
|
|
matplotlib
|
|
numpy
|
|
opencv4
|
|
pandas
|
|
pillow
|
|
plotly
|
|
scipy
|
|
scikit-image
|
|
scikit-learn
|
|
seaborn
|
|
tokenizers
|
|
transformers
|
|
];
|
|
linuxOnlyPackages = lib.optionals pkgs.stdenv.isLinux (
|
|
with ps;
|
|
[
|
|
sentencepiece
|
|
torch
|
|
torchaudio
|
|
torchvision
|
|
]
|
|
);
|
|
in
|
|
basePackages ++ linuxOnlyPackages
|
|
);
|
|
in
|
|
pkgs.mkShell {
|
|
name = "ai-notebook";
|
|
packages = with pkgs; [
|
|
pythonWithAI
|
|
uv
|
|
git-lfs
|
|
ruff
|
|
pyright
|
|
];
|
|
|
|
shellHook = ''
|
|
export UV_PYTHON="${pythonWithAI}/bin/python3"
|
|
export UV_LINK_MODE=copy
|
|
|
|
export PYTHONNOUSERSITE=1
|
|
export HF_HOME="''${XDG_CACHE_HOME:-$HOME/.cache}/huggingface"
|
|
export HF_HUB_CACHE="$HF_HOME/hub"
|
|
export TRANSFORMERS_CACHE="$HF_HOME/transformers"
|
|
export JUPYTER_CONFIG_DIR="''${XDG_CONFIG_HOME:-$HOME/.config}/jupyter"
|
|
|
|
mkdir -p "$HF_HOME" "$JUPYTER_CONFIG_DIR"
|
|
|
|
echo "AI notebook shell ready"
|
|
echo "Python: ${python.version}"
|
|
echo "HuggingFace cache: $HF_HOME"
|
|
echo "Launch JupyterLab: jupyter lab"
|
|
echo "Launch classic notebook: jupyter notebook"
|
|
'';
|
|
}
|