pass run_llm_in_parallel to benchmarks

Reviewed By: SimonWan Differential Revision: D69662273 fbshipit-source-id: 93e109e5c851cbdb3f6772d93f76936169184630
meta-llama · Feb 14, 2025 · 23156b7 · 23156b7
1 parent c6dee62
commit 23156b7
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/CybersecurityBenchmarks/benchmark/benchmark.py b/CybersecurityBenchmarks/benchmark/benchmark.py
@@ -94,6 +94,7 @@ def query_llm_to_generate_responses(
     async def run(
         self,
         num_test_cases: int = 0,
+        run_llm_in_parallel: int = 1,
     ) -> None:
         """
         Processes a LLM-response.

diff --git a/CybersecurityBenchmarks/benchmark/instruct_or_autocomplete_benchmark.py b/CybersecurityBenchmarks/benchmark/instruct_or_autocomplete_benchmark.py
@@ -49,7 +49,7 @@ def query_llm_to_generate_responses(
 
         self.response_path.write_text(json.dumps(response_result, indent=4))
 
-    async def run(self, num_test_cases: int = 0) -> None:  # noqa: C901
+    async def run(self, num_test_cases: int = 0, run_llm_in_parallel: int = 1) -> None:  # noqa: C901
         """
         Processes a LLM-response by checking whether the response is detected as insecure by ICD.
         """

diff --git a/CybersecurityBenchmarks/benchmark/run.py b/CybersecurityBenchmarks/benchmark/run.py
@@ -280,7 +280,7 @@ def main(
                 )
             if args.num_queries_per_prompt:
                 args.num_test_cases = args.num_test_cases * args.num_queries_per_prompt
-        asyncio.run(benchmark.run(args.num_test_cases))
+        asyncio.run(benchmark.run(args.num_test_cases, args.run_llm_in_parallel))
 
 
 if __name__ == "__main__":