pass device to transfomers pipeline

When a `pipeline` is created the base class attempts to auto-detect the optimal hardware, since the project accepts configuration for hardware device selection the device must be passed. Signed-off-by: Jeffrey Martin <[email protected]>
NVIDIA · Dec 11, 2024 · 307fe77 · 307fe77
1 parent 0b837c1
commit 307fe77
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 3 deletions.
diff --git a/garak/detectors/base.py b/garak/detectors/base.py
@@ -120,7 +120,9 @@ def __init__(self, config_root=_config):
             self.detector_model_path
         )
         self.detector = TextClassificationPipeline(
-            model=self.detector_model, tokenizer=self.detector_tokenizer
+            model=self.detector_model,
+            tokenizer=self.detector_tokenizer,
+            device=self.device,
         )
 
         transformers_logging.set_verbosity(orig_loglevel)

diff --git a/pyproject.toml b/pyproject.toml
@@ -45,7 +45,7 @@ classifiers = [
 requires-python = ">=3.10"
 dependencies = [
   "base2048>=0.1.3",
-  "transformers>=4.43.0,<4.47.0",
+  "transformers>=4.43.0",
   "datasets>=2.14.6,<2.17",
   "colorama>=0.4.3",
   "tqdm>=4.64.0",

diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,5 @@
 base2048>=0.1.3
-transformers>=4.43.0,<4.47.0
+transformers>=4.43.0
 datasets>=2.14.6,<2.17
 colorama>=0.4.3
 tqdm>=4.64.0