nix-ai-template/flake.nix
2024-07-11 23:55:18 +02:00

55 lines
1.2 KiB
Nix

{
description = "My AI Project";
inputs = {
nix-ai.url = "git+https://git.wavelens.io/public/nix-ai";
};
outputs = { nix-ai, ... }: nix-ai.lib.mkFlake {
presets.jupyter = true;
packages = [ "jq" ];
pythonPackages = [ "tqdm" ];
environmentVariables = {
HF_HOME = ".cache/huggingface";
};
datasets = {
mnist = {
src = {
url = "https://static.wavelens.io/private/corpus2.pkl";
hash = "sha256-RKOl/I/J1m8uKDm+XdVT9NcBiJspHM4m4lKIJ0aObJE=";
};
prepare = {
directoryPath = ./utils;
drop = [ "train-images-idx3-ubyte" ];
commands = ''
python clean.py
'';
};
};
};
trainings = {
mnist-ai = {
gpu = "H100-1";
directoryPath = ./minst-ai;
copyDatasets = [ "mnist" ];
drop = [ "model.pt" "runs" "test.txt" ];
commands = ''
mkdir -p temp
python train.py
python test.py | tee test.txt
'';
configurations = [{
model = "minst-model";
dataset = "mnist";
layersA = [5 10];
layersB = [5 10];
epochs = 1;
}];
};
};
};
}