yujunjun/haystack
1
0
Fork 0
mirror of https://github.com/deepset-ai/haystack.git synced 2025-11-02 10:49:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: Change use_auth_token to token in TransformersQueryClassifier (#5659)

Browse Source
This commit is contained in:
bogdankostic 2023-08-29 15:21:25 +02:00 committed by GitHub
parent ee13125e06
commit 07c85905f3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 10 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
haystack/nodes/query_classifier/transformers.py
View File

@ -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

9
test/nodes/test_query_classifier.py
View File

@ -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(