diff --git a/haystack/nodes/query_classifier/transformers.py b/haystack/nodes/query_classifier/transformers.py index 4a4f1e01b..7fb926e04 100644 --- a/haystack/nodes/query_classifier/transformers.py +++ b/haystack/nodes/query_classifier/transformers.py @@ -121,7 +121,7 @@ class TransformersQueryClassifier(BaseQueryClassifier): tokenizer=tokenizer, device=resolved_devices[0], revision=model_version, - use_auth_token=use_auth_token, + token=use_auth_token, ) self.labels = labels diff --git a/test/nodes/test_query_classifier.py b/test/nodes/test_query_classifier.py index acc8907e7..50021f635 100644 --- a/test/nodes/test_query_classifier.py +++ b/test/nodes/test_query_classifier.py @@ -1,3 +1,4 @@ +from unittest.mock import patch import pytest from pathlib import Path from urllib.error import URLError @@ -15,6 +16,14 @@ def test_sklearnqueryclassifier_deprecation(): pass +@pytest.mark.unit +def test_query_classifier_initialized_with_token_instead_of_use_auth_token(): + with patch("haystack.nodes.query_classifier.transformers.pipeline") as mock_transformers_pipeline: + classifier = TransformersQueryClassifier(task="zero-shot-classification") + assert "token" in mock_transformers_pipeline.call_args.kwargs + assert "use_auth_token" not in mock_transformers_pipeline.call_args.kwargs + + @pytest.fixture def transformers_query_classifier(): return TransformersQueryClassifier(