Format Benchmark Code (#399)

benchmark/long_json_decode/bench_sglang.py

@@ -2,10 +2,12 @@ import argparse
 import json
 import time
 
-import numpy as np
 import sglang as sgl
-from sglang.test.test_utils import add_common_sglang_args_and_parse, select_sglang_backend
-from sglang.utils import read_jsonl, dump_state_text
+from sglang.test.test_utils import (
+    add_common_sglang_args_and_parse,
+    select_sglang_backend,
+)
+from sglang.utils import dump_state_text, read_jsonl
 
 
 @sgl.function
@@ -13,21 +15,31 @@ def json_decode(s, document):
     s += "Please extract the information of a city from the following wikipedia page.\n"
     s += "Page begin.\n" + document + "Page end.\n"
     s += "Here is the name, country, and symbol of the city in JSON format.\n"
-    s += '{\n'
+    s += "{\n"
     s += '  "name": "' + sgl.gen("name", max_tokens=8, stop='"') + '",\n'
     s += '  "country": "' + sgl.gen("country", max_tokens=8, stop='"') + '",\n'
-    s += '  "air port code": "' + sgl.gen("air port code", max_tokens=8, stop='"') + '",\n'
-    s += '  "top 3 landmarks": "' + sgl.gen("landmarks", max_tokens=24, stop='"') + '",\n'
-    s += '}\n'
+    s += (
+        '  "air port code": "'
+        + sgl.gen("air port code", max_tokens=8, stop='"')
+        + '",\n'
+    )
+    s += (
+        '  "top 3 landmarks": "'
+        + sgl.gen("landmarks", max_tokens=24, stop='"')
+        + '",\n'
+    )
+    s += "}\n"
 
 
 def main(args):
     lines = read_jsonl(args.data_path)
     arguments = []
-    for i in range(len(lines[:args.num_questions])):
-        arguments.append({
-            "document": lines[i]["document"],
-        })
+    for i in range(len(lines[: args.num_questions])):
+        arguments.append(
+            {
+                "document": lines[i]["document"],
+            }
+        )
 
     # Select backend
     backend = select_sglang_backend(args)
@@ -36,10 +48,11 @@ def main(args):
     # Run requests
     tic = time.time()
     states = json_decode.run_batch(
-        arguments, temperature=0, num_threads=args.parallel, progress_bar=True)
+        arguments, temperature=0, num_threads=args.parallel, progress_bar=True
+    )
     latency = time.time() - tic
 
-    # Compute accuracy 
+    # Compute accuracy
     print(f"Latency: {latency:.3f}")
 
     # Write results
@@ -55,7 +68,7 @@ def main(args):
             "other": {
                 "num_questions": args.num_questions,
                 "parallel": args.parallel,
-            }
+            },
         }
         fout.write(json.dumps(value) + "\n")