From a8c22c11f4c5daca60aed33fee6db38c578ce19e Mon Sep 17 00:00:00 2001
From: David Baynard <git@baynard.dev>
Date: Wed, 14 Aug 2024 02:41:22 +0100
Subject: [PATCH] Add huggingface fetcher

---
 utils/fetchers/default.nix              | 32 ++++++++++-----
 utils/fetchers/fetchFromHuggingFace.nix | 53 +++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 10 deletions(-)
 create mode 100644 utils/fetchers/fetchFromHuggingFace.nix

diff --git a/utils/fetchers/default.nix b/utils/fetchers/default.nix
index bcf0f88..96c7147 100644
--- a/utils/fetchers/default.nix
+++ b/utils/fetchers/default.nix
@@ -6,16 +6,28 @@ in
 {
   overlays.fetchers = final: prev:
     let
+      fetchFromHuggingFace = prev.callPackage ./fetchFromHuggingFace.nix {
+        inherit (final) huggingface-model-downloader;
+      };
       extras = { };
     in
-    lib.foldFor pnames (pname: {
-      ${pname} = prev.callPackage
-        (./. + "/${pname}.nix")
-        (extras."${pname}" or { });
-    });
+    lib.foldFor pnames
+      (pname: {
+        ${pname} = prev.callPackage
+          (./. + "/${pname}.nix")
+          (extras."${pname}" or { });
+      }) // {
+      inherit fetchFromHuggingFace;
+    };
 } //
-lib.foldFor lib.platforms.all (system: {
-  packages.${system} = self.overlays.fetchers
-    (nixpkgs.legacyPackages.${system} // self.packages.${system})
-    nixpkgs.legacyPackages.${system};
-})
+lib.foldFor lib.platforms.all (system:
+  let
+    pkgs = nixpkgs.legacyPackages.${system};
+  in
+  {
+    packages.${system} =
+      lib.filterAttrs (_: lib.isDerivation) self.legacyPackages.${system};
+    legacyPackages.${system} = self.overlays.fetchers
+      (pkgs // self.packages.${system})
+      pkgs;
+  })
diff --git a/utils/fetchers/fetchFromHuggingFace.nix b/utils/fetchers/fetchFromHuggingFace.nix
new file mode 100644
index 0000000..8fc4a01
--- /dev/null
+++ b/utils/fetchers/fetchFromHuggingFace.nix
@@ -0,0 +1,53 @@
+{ lib
+, huggingface-model-downloader
+, stdenvNoCC
+}:
+
+{ model
+, version
+, hash ? null
+, requireToken ? false
+, extraArgs ? [ ]
+}:
+
+stdenvNoCC.mkDerivation {
+  pname = "hf-${lib.replaceStrings ["/"] ["-"] model}";
+  inherit version;
+
+  dontUnpack = true;
+  dontFixup = true;
+
+  impureEnvVars = [ "HUGGING_FACE_HUB_TOKEN" ];
+
+  nativeBuildInputs = [
+    huggingface-model-downloader
+  ];
+
+  preBuild = lib.optionalString requireToken ''
+    if [ -z ''${HUGGING_FACE_HUB_TOKEN} ]; then
+      echo "Private: fetchFromHuggingFace requires the nix building process (nix-daemon in multi user mode) to have the HUGGING_FACE_HUB_TOKEN env var set." >&2
+      exit 1
+    fi
+  '';
+
+  EXTRA_ARGS = extraArgs;
+
+  buildPhase = ''
+    runHook preBuild
+
+    ${lib.getExe huggingface-model-downloader} -m ${model} "''${EXTRA_ARGS[@]}"
+
+    runHook postBuild
+  '';
+
+  installPhase = ''
+    runHook preInstall
+
+    cp -r --reflink=auto . "$out"
+
+    runHook postInstall
+  '';
+
+  outputHash = hash;
+  outputHashMode = "recursive";
+}