sticker2-python/default.nix at master · stickeritis/sticker2-python · 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
{
  pkgs ? import (import ./nix/sources.nix).nixpkgs {
    config = {
      allowUnfreePredicate = pkg: builtins.elem (pkgs.lib.getName pkg) [
        "libtorch"
      ];
    };
  }

, python3Packages ? pkgs.python3Packages
}:

let
  sources = import ./nix/sources.nix;
  libtorch = pkgs.libtorch-bin;
  crateOverrides = with pkgs; defaultCrateOverrides // {
    pyo3 = attr: {
      buildInputs = [ python3Packages.python ];
    };

    sentencepiece-sys = attr: {
      nativeBuildInputs = [ pkgconfig ];

      buildInputs = [ sentencepiece ];
    };

    sticker2-python = attr: rec {
      pname = "sticker2-python";
      name = "${pname}-${attr.version}";

      buildInputs = [ libtorch ] ++
        lib.optional stdenv.isDarwin darwin.Security;

      features = [];

      installPhase = let
        sitePackages = python3Packages.python.sitePackages;
        sharedLibrary = stdenv.hostPlatform.extensions.sharedLibrary;
      in ''
        runHook preInstall

        mkdir -p "$out/${sitePackages}"
        cp target/lib/libsticker2_python*${sharedLibrary} \
          "$out/${sitePackages}/sticker2.so"
        export PYTHONPATH="$out/${sitePackages}:$PYTHONPATH"

        runHook postInstall
      '';
    };

    torch-sys = attr: {
      buildInputs = lib.optional stdenv.isDarwin curl;

      LIBTORCH = "${libtorch.dev}";
    };
  };
  buildRustCrate = pkgs.buildRustCrate.override {
    defaultCrateOverrides = crateOverrides;
  };
  crateTools = pkgs.callPackage "${sources.crate2nix}/tools.nix" {};
  cargoNix = pkgs.callPackage (crateTools.generatedCargoNix {
    name = "sticker2";
    src = pkgs.nix-gitignore.gitignoreSource [ ".git/" "nix/" "*.nix" ] ./.;
  }) {
    inherit buildRustCrate;
  };
in cargoNix.rootCrate.build