python312Packages.colbert-ai: fix import

2025-03-24 23:20:17 +01:00
parent 8161b1cb1e
commit f96eda85f3
1 changed files with 9 additions and 0 deletions
@@ -32,6 +32,15 @@ buildPythonPackage rec {
    hash = "sha256-qNb9tOInLysI7Tf45QlgchYNhBXR5AWFdRiYt35iW6s=";
  };

+  # ImportError: cannot import name 'AdamW' from 'transformers'
+  # https://github.com/stanford-futuredata/ColBERT/pull/390
+  postPatch = ''
+    substituteInPlace colbert/training/training.py \
+      --replace-fail \
+      "from transformers import AdamW, get_linear_schedule_with_warmup" \
+      "from transformers import get_linear_schedule_with_warmup; from torch.optim import AdamW"
+  '';
+
  pythonRemoveDeps = [ "git-python" ];

  build-system = [