This commit is contained in:
Ashwin Bharambe 2025-01-22 18:09:46 -08:00
parent a8345f5f76
commit 72a1b27d01

View file

@ -7,7 +7,7 @@
def pytest_addoption(parser): def pytest_addoption(parser):
parser.addoption( parser.addoption(
"--safety_shield", "--safety-shield",
action="store", action="store",
default="meta-llama/Llama-Guard-3-1B", default="meta-llama/Llama-Guard-3-1B",
help="Specify the safety shield model to use for testing", help="Specify the safety shield model to use for testing",
@ -18,5 +18,5 @@ def pytest_generate_tests(metafunc):
if "llama_guard_text_shield_id" in metafunc.fixturenames: if "llama_guard_text_shield_id" in metafunc.fixturenames:
metafunc.parametrize( metafunc.parametrize(
"llama_guard_text_shield_id", "llama_guard_text_shield_id",
[metafunc.config.getoption("--safety_shield")], [metafunc.config.getoption("--safety-shield")],
) )