Skip to content

Commit 6b1391c

Browse files
authored
[Misc] refactor neuron_multimodal and profiling (vllm-project#19397)
Signed-off-by: reidliu41 <reid201711@gmail.com> Co-authored-by: reidliu41 <reid201711@gmail.com>
1 parent a3f66e7 commit 6b1391c

File tree

2 files changed

+11
-3
lines changed

2 files changed

+11
-3
lines changed

examples/offline_inference/neuron_multimodal.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ def print_outputs(outputs):
6464
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
6565

6666

67-
if __name__ == "__main__":
67+
def main():
6868
assert (
6969
len(PROMPTS) == len(IMAGES) == len(SAMPLING_PARAMS)
7070
), f"""Text, image prompts and sampling parameters should have the
@@ -104,3 +104,7 @@ def print_outputs(outputs):
104104
# test batch-size = 4
105105
outputs = llm.generate(batched_inputs, batched_sample_params)
106106
print_outputs(outputs)
107+
108+
109+
if __name__ == "__main__":
110+
main()

examples/offline_inference/profiling_tpu/profiling.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ def run_to_completion():
7070
return
7171

7272

73-
if __name__ == "__main__":
73+
def parse_args():
7474
parser = FlexibleArgumentParser(
7575
description="Benchmark the latency of processing a single batch of "
7676
"requests till completion."
@@ -102,5 +102,9 @@ def run_to_completion():
102102
)
103103

104104
parser = EngineArgs.add_cli_args(parser)
105-
args = parser.parse_args()
105+
return parser.parse_args()
106+
107+
108+
if __name__ == "__main__":
109+
args = parse_args()
106110
main(args)

0 commit comments

Comments
 (0)