Skip to content

Commit

Permalink
pass device to transfomers pipeline (#1042)
Browse files Browse the repository at this point in the history
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.

## Verification

- [x] actions macOS test automation
  • Loading branch information
leondz authored Dec 12, 2024
2 parents 0b837c1 + 307fe77 commit 4305c5a
Show file tree
Hide file tree
Showing 3 changed files with 5 additions and 3 deletions.
4 changes: 3 additions & 1 deletion garak/detectors/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
2 changes: 1 addition & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
@@ -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
Expand Down

0 comments on commit 4305c5a

Please sign in to comment.