53 lines
1.4 KiB
Nix
53 lines
1.4 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"
|
|
'';
|
|
}
|