From 94304649f1d2e7eb9ff2b633c592eeb7a4343fc9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Thu, 12 Sep 2024 10:44:01 +0200
Subject: [PATCH] nix: support Python tokenizer conversion in the router
 (#2515)

Ideally we wouldn't have the router wrapper that this change adds,
but when I give PyO3 a Python interpreter with packages, it ends
up linking libpython from the Python interpreter rather than the
constructed environment and cannot pick up the Python modules as
a result.
---
 flake.nix               | 23 ++++++++++++++++++++---
 nix/crate-overrides.nix |  3 +++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/flake.nix b/flake.nix
index 58a8e311..5be413c5 100644
--- a/flake.nix
+++ b/flake.nix
@@ -46,9 +46,26 @@
         launcher = cargoNix.workspaceMembers.text-generation-launcher.build.override {
           inherit crateOverrides;
         };
-        router = cargoNix.workspaceMembers.text-generation-router-v3.build.override {
-          inherit crateOverrides;
-        };
+        router =
+          let
+            routerUnwrapped = cargoNix.workspaceMembers.text-generation-router-v3.build.override {
+              inherit crateOverrides;
+            };
+            packagePath =
+              with pkgs.python3.pkgs;
+              makePythonPath [
+                protobuf
+                sentencepiece
+                torch
+                transformers
+              ];
+          in
+          pkgs.writeShellApplication {
+            name = "text-generation-router";
+            text = ''
+              PYTHONPATH="${packagePath}" ${routerUnwrapped}/bin/text-generation-router "$@"
+            '';
+          };
         server = pkgs.python3.pkgs.callPackage ./nix/server.nix { inherit nix-filter; };
       in
       {
diff --git a/nix/crate-overrides.nix b/nix/crate-overrides.nix
index cbf366af..a4e74c6d 100644
--- a/nix/crate-overrides.nix
+++ b/nix/crate-overrides.nix
@@ -28,6 +28,9 @@ defaultCrateOverrides
       ];
     };
   };
+  pyo3-build-config = attrs: {
+    buildInputs = [ python3 ];
+  };
   text-generation-benchmark = attrs: {
     src = filter {
       root = ../benchmark;