Skip to content

Commit

Permalink
doc string
Browse files Browse the repository at this point in the history
  • Loading branch information
horheynm committed Jan 5, 2024
1 parent 133d6ca commit a797562
Showing 1 changed file with 23 additions and 0 deletions.
23 changes: 23 additions & 0 deletions src/deepsparse/benchmark/benchmark_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -348,6 +348,29 @@ def benchmark_model(
export_path: Optional[str] = None,
disable_kv_cache_overrides: bool = False,
) -> Dict:
"""
Run benchamrk of a given model on a given engine
:param model_path: Path or stub to the model
:param batch_size: Batch size, Defaults to 1.
:param input_shapes: Shape of the input. Defaults to "".
:param num_cores: Number of cores to use. Defaults to None.
:param scenario: The "type" of scenario to use. Defaults to "sync".
:param time: Total run-time to run the benchmark for. Defaults to 10.
:param warmup_time: Time to run the engine before running benchmarks. Defaults to 2.
:param num_streams: Number of streams. Defaults to None.
:param sequence_length: Sequence length for text-gen. Defaults to None.
:param input_ids_length: Length of the input ids. Defaults to 1.
:param thread_pinning: The hardware to pin for threading. Defaults to "core".
:param engine: The type of engine to use. Defaults to our proprietary DEEPSPARSE_ENGINE.
:param internal_kv_cache: Runs the benchmark with or w/o the kv cache. Defaults to True.
:param quiet: Verbose option if set to False. Defaults to False.
:param export_path: Path to save the results. Defaults to None.
:param disable_kv_cache_overrides: Option to override the kv cache. Defaults to False.
:returns: Dictionary of benchmarked metrics
"""

if quiet:
set_logging_level(logging.WARN)

Expand Down

0 comments on commit a797562

Please sign in to comment.