diff --git a/garak/generators/huggingface.py b/garak/generators/huggingface.py index 86a94bee..d1418986 100644 --- a/garak/generators/huggingface.py +++ b/garak/generators/huggingface.py @@ -256,7 +256,7 @@ def __init__(self, name="", config_root=_config): self.name = name super().__init__(self.name, config_root=config_root) - self.uri = self.URI + name + self.uri = self.URI + self.name # special case for api token requirement this also reserves `headers` as not configurable if self.api_key: @@ -376,7 +376,7 @@ class InferenceEndpoint(InferenceAPI): def __init__(self, name="", config_root=_config): super().__init__(name, config_root=config_root) - self.uri = name + self.uri = self.name @backoff.on_exception( backoff.fibo, diff --git a/tests/generators/conftest.py b/tests/generators/conftest.py index bea801f9..9a760d80 100644 --- a/tests/generators/conftest.py +++ b/tests/generators/conftest.py @@ -11,3 +11,10 @@ def openai_compat_mocks(): """Mock responses for OpenAI compatible endpoints""" with open(pathlib.Path(__file__).parents[0] / "openai.json") as mock_openai: return json.load(mock_openai) + + +@pytest.fixture +def hf_endpoint_mocks(): + """Mock responses for Huggingface InferenceAPI based endpoints""" + with open(pathlib.Path(__file__).parents[0] / "hf_inference.json") as mock_openai: + return json.load(mock_openai) diff --git a/tests/generators/hf_inference.json b/tests/generators/hf_inference.json new file mode 100644 index 00000000..9cd1ddcf --- /dev/null +++ b/tests/generators/hf_inference.json @@ -0,0 +1,10 @@ +{ + "hf_inference": { + "code": 200, + "json": [ + { + "generated_text":"restricted by their policy," + } + ] + } +} diff --git a/tests/generators/test_huggingface.py b/tests/generators/test_huggingface.py index 54491db7..f784d95d 100644 --- a/tests/generators/test_huggingface.py +++ b/tests/generators/test_huggingface.py @@ -1,4 +1,5 @@ import pytest +import requests import transformers import garak.generators.huggingface from garak._config import GarakSubConfig @@ -8,6 +9,7 @@ def hf_generator_config(): gen_config = { "huggingface": { + "api_key": "fake", "hf_args": { "device": "cpu", "torch_dtype": "float32", @@ -19,6 +21,17 @@ def hf_generator_config(): return config_root +@pytest.fixture +def hf_mock_response(hf_endpoint_mocks): + import json + + mock_resp_data = hf_endpoint_mocks["hf_inference"] + mock_resp = requests.Response() + mock_resp.status_code = mock_resp_data["code"] + mock_resp._content = json.dumps(mock_resp_data["json"]).encode("UTF-8") + return mock_resp + + def test_pipeline(hf_generator_config): generations = 10 g = garak.generators.huggingface.Pipeline("gpt2", config_root=hf_generator_config) @@ -37,16 +50,55 @@ def test_pipeline(hf_generator_config): assert isinstance(item, str) -def test_inference(): - return # slow w/o key - g = garak.generators.huggingface.InferenceAPI("gpt2") - assert g.name == "gpt2" +def test_inference(mocker, hf_mock_response, hf_generator_config): + model_name = "gpt2" + mock_request = mocker.patch.object( + requests, "request", return_value=hf_mock_response + ) + + g = garak.generators.huggingface.InferenceAPI( + model_name, config_root=hf_generator_config + ) + assert g.name == model_name + assert model_name in g.uri + + hf_generator_config.generators["huggingface"]["name"] = model_name + g = garak.generators.huggingface.InferenceAPI(config_root=hf_generator_config) + assert g.name == model_name + assert model_name in g.uri + assert isinstance(g.max_tokens, int) + g.max_tokens = 99 + assert g.max_tokens == 99 + g.temperature = 0.1 + assert g.temperature == 0.1 + output = g.generate("") + mock_request.assert_called_once() + assert len(output) == 1 # 1 generation by default + for item in output: + assert isinstance(item, str) + + +def test_endpoint(mocker, hf_mock_response, hf_generator_config): + model_name = "https://localhost:8000/gpt2" + mock_request = mocker.patch.object(requests, "post", return_value=hf_mock_response) + + g = garak.generators.huggingface.InferenceEndpoint( + model_name, config_root=hf_generator_config + ) + assert g.name == model_name + assert g.uri == model_name + + hf_generator_config.generators["huggingface"]["name"] = model_name + g = garak.generators.huggingface.InferenceEndpoint(config_root=hf_generator_config) + assert g.name == model_name + assert g.uri == model_name assert isinstance(g.max_tokens, int) g.max_tokens = 99 assert g.max_tokens == 99 g.temperature = 0.1 assert g.temperature == 0.1 output = g.generate("") + mock_request.assert_called_once() assert len(output) == 1 # 1 generation by default for item in output: assert isinstance(item, str)