From 4afc39c5f8d25812dd677a3f9319c890a18ba9bc Mon Sep 17 00:00:00 2001 From: reidliu41 Date: Tue, 10 Jun 2025 12:33:49 +0800 Subject: [PATCH] [Misc] refactor neuron_multimodal and profiling Signed-off-by: reidliu41 --- examples/offline_inference/neuron_multimodal.py | 6 +++++- examples/offline_inference/profiling_tpu/profiling.py | 8 ++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/examples/offline_inference/neuron_multimodal.py b/examples/offline_inference/neuron_multimodal.py index 6ff8faabd748..26f7505f2fa5 100644 --- a/examples/offline_inference/neuron_multimodal.py +++ b/examples/offline_inference/neuron_multimodal.py @@ -64,7 +64,7 @@ def print_outputs(outputs): print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") -if __name__ == "__main__": +def main(): assert ( len(PROMPTS) == len(IMAGES) == len(SAMPLING_PARAMS) ), f"""Text, image prompts and sampling parameters should have the @@ -104,3 +104,7 @@ def print_outputs(outputs): # test batch-size = 4 outputs = llm.generate(batched_inputs, batched_sample_params) print_outputs(outputs) + + +if __name__ == "__main__": + main() diff --git a/examples/offline_inference/profiling_tpu/profiling.py b/examples/offline_inference/profiling_tpu/profiling.py index 5200be82694a..dfcbd8c8d360 100644 --- a/examples/offline_inference/profiling_tpu/profiling.py +++ b/examples/offline_inference/profiling_tpu/profiling.py @@ -70,7 +70,7 @@ def run_to_completion(): return -if __name__ == "__main__": +def parse_args(): parser = FlexibleArgumentParser( description="Benchmark the latency of processing a single batch of " "requests till completion." @@ -102,5 +102,9 @@ def run_to_completion(): ) parser = EngineArgs.add_cli_args(parser) - args = parser.parse_args() + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() main(args)