peft/shell.nix at main · adconner/peft · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
let
  # # https://status.nixos.org/ nixos-unstable-small 8/5/25
  # # https://hydra.nix-community.org/job/nixpkgs/cuda/python3Packages.torch.x86_64-linux
  # pkgs = import (fetchTarball "https://github.com/NixOS/nixpkgs/tarball/9ebe222ec7ef9de52478f76cba3f0324c1d1119f") {
  #   config.allowUnfree = true;
  #   config.cudaSupport = true;
  # };
  pkgs = import (fetchTarball "https://github.com/NixOS/nixpkgs/tarball/e9f00bd893984bc8ce46c895c3bf7cac95331127") {
    overlays = [
      (self: super: {
        pythonPackagesExtensions = super.pythonPackagesExtensions ++ [
          (python-self: python-super: {
            flax = python-super.flax.overridePythonAttrs (_: { doCheck = false; });
            torch = python-super.torch.override { cudaSupport = true; };

            # aqtp = python-self.buildPythonPackage rec {
            #   pname = "aqtp";
            #   version = "0.9.0";
            #   src = pkgs.fetchFromGitHub {
            #     owner = "google";
            #     repo = "aqt";
            #     rev = "e781c13f63d359027922b31e7e227cca4fab4dc8";
            #     sha256 = "sha256-FvqMD6OLvcK1umCRW8skAMjUHw5koy/a7xIn76mG9wU=";
            #   };
            #   pyproject = true;
            #   build-system = with python-self; [ setuptools ];
            #   propagatedBuildInputs = with python-self; [ jax jaxlib absl-py flax ];
            #   doCheck = false;
            #   patchPhase = "cat setup.py | sed 's/0.1.0/${version}/' > tmp && mv tmp setup.py";
            # };

            # haliax = let
            # in python-self.buildPythonPackage rec {
            #   pname = "haliax";
            #   version = "";
            #   src = pkgs.fetchFromGitHub {
            #     owner = "stanford-crfm";
            #     repo = pname;
            #     rev = "30066d0a1eea3d282580a5a5fef815211d005abb";
            #     sha256 = "sha256-/gzDORDv3GWWxMxQO7OE7C+9kn0v2dE+ojvdvhy/pKk=";
            #   };
            #   pyproject = true;
            #   build-system = with python-self; [ hatchling ];
            #   propagatedBuildInputs = with python-self; [ equinox jaxtyping jmp safetensors aqtp ];
            # };

            torch2jax = pythonPackages.buildPythonPackage rec {
              pname = "torch2jax";
              version = "0.1.0";
              src = ../torch2jax;
              # src = pkgs.fetchFromGitHub {
              #   owner = "samuela";
              #   repo = pname;
              #   rev = "505769d32fc20b95e23f53c36dd320db31066282";
              #   sha256 = "sha256-+1AICchWfmRou93iMq53Ai+1KkxYS8rBU14A9pjt3a0=";
              # };
              pyproject = true;
              build-system = with pythonPackages; [ setuptools ];
              propagatedBuildInputs = with pythonPackages; [ torch jax ];
            };

          })
        ];
      })
    ];
  };
  pythonPackages = pkgs.python3Packages;
in
pkgs.mkShell {
  buildInputs = with pythonPackages; [
    jax-cuda12-plugin
    jax
    optax
    equinox
    jaxtyping
    chex
    flax
    torch2jax
    draccus

    datasets
    # evaluate
    # scikit-learn
    bokeh

    accelerate
    transformers
    torch

    ipython
    ruff
    venvShellHook
  ] ++ (with pkgs; [
    pyright
    pandoc
  ]);
  venvDir = ".venv";
}