diff --git a/garak/probes/fileformats.py b/garak/probes/fileformats.py index 6c87e970..e4542b52 100644 --- a/garak/probes/fileformats.py +++ b/garak/probes/fileformats.py @@ -14,7 +14,7 @@ from garak.configurable import Configurable from garak.probes.base import Probe import garak.attempt - +import garak.resources.theme class HF_Files(Probe, Configurable): """Get a manifest of files associated with a Hugging Face generator @@ -50,7 +50,7 @@ def probe(self, generator) -> Iterable[garak.attempt.Attempt]: return [] if generator.__class__.__name__ not in self.supported_generators: return [] - attempt = self._mint_attempt(None) + attempt = self._mint_attempt(generator.name) repo_filenames = huggingface_hub.list_repo_files(generator.name) local_filenames = [] @@ -58,7 +58,7 @@ def probe(self, generator) -> Iterable[garak.attempt.Attempt]: repo_filenames, leave=False, desc=f"Gathering files in {generator.name}", - colour="#e5a70e", # probe theme colour + colour=f"#{garak.resources.theme.PROBE_RGB}", ): local_filename = huggingface_hub.hf_hub_download( generator.name, repo_filename, force_download=False diff --git a/tests/probes/test_probes_fileformats.py b/tests/probes/test_probes_fileformats.py new file mode 100644 index 00000000..16ba21ed --- /dev/null +++ b/tests/probes/test_probes_fileformats.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os + +import garak._config +import garak._plugins + +import garak.probes.base +import garak.probes.fileformats +import garak.attempt + + +def test_hf_files_load(): + p = garak.probes.fileformats.HF_Files() + assert isinstance(p, garak.probes.base.Probe) + + +def test_hf_files_hf_repo(): + p = garak._plugins.load_plugin("probes.fileformats.HF_Files") + garak._config.plugins.generators["huggingface"] = {"Model": {"name": "gpt2"}} + g = garak._plugins.load_plugin( + "generators.huggingface.Model", config_root=garak._config + ) + r = p.probe(g) + assert isinstance(r, list), ".probe should return a list" + assert len(r) == 1, "HF_Files.probe() should return one attempt" + assert isinstance( + r[0], garak.attempt.Attempt + ), "HF_Files.probe() must return an Attempt" + assert isinstance(r[0].outputs, list), "File list scan should return a list" + assert len(r[0].outputs) > 0, "File list scan should return list of filenames" + for filename in r[0].outputs: + assert isinstance( + filename, str + ), "File list scan should return list of string filenames" + assert os.path.isfile(filename), "List of HF_Files paths should all be real"