File tree Expand file tree Collapse file tree 2 files changed +11
-3
lines changed
examples/offline_inference Expand file tree Collapse file tree 2 files changed +11
-3
lines changed Original file line number Diff line number Diff line change @@ -64,7 +64,7 @@ def print_outputs(outputs):
64
64
print (f"Prompt: { prompt !r} , Generated text: { generated_text !r} " )
65
65
66
66
67
- if __name__ == "__main__" :
67
+ def main () :
68
68
assert (
69
69
len (PROMPTS ) == len (IMAGES ) == len (SAMPLING_PARAMS )
70
70
), f"""Text, image prompts and sampling parameters should have the
@@ -104,3 +104,7 @@ def print_outputs(outputs):
104
104
# test batch-size = 4
105
105
outputs = llm .generate (batched_inputs , batched_sample_params )
106
106
print_outputs (outputs )
107
+
108
+
109
+ if __name__ == "__main__" :
110
+ main ()
Original file line number Diff line number Diff line change @@ -70,7 +70,7 @@ def run_to_completion():
70
70
return
71
71
72
72
73
- if __name__ == "__main__" :
73
+ def parse_args () :
74
74
parser = FlexibleArgumentParser (
75
75
description = "Benchmark the latency of processing a single batch of "
76
76
"requests till completion."
@@ -102,5 +102,9 @@ def run_to_completion():
102
102
)
103
103
104
104
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 ()
106
110
main (args )
You can’t perform that action at this time.
0 commit comments