Format Benchmark Code (#399)

benchmark/long_json_decode/bench_other.py

@@ -1,21 +1,25 @@
 import argparse
-import asyncio
-from concurrent.futures import ThreadPoolExecutor
-from functools import partial
 import json
 import time
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
 
 from tqdm import tqdm
-import numpy as np
-from sglang.test.test_utils import add_common_other_args_and_parse, call_generate_lightllm, call_generate_vllm, call_generate_srt_raw
-from sglang.utils import read_jsonl, dump_state_text
+
+from sglang.test.test_utils import (
+    add_common_other_args_and_parse,
+    call_generate_lightllm,
+    call_generate_srt_raw,
+    call_generate_vllm,
+)
+from sglang.utils import dump_state_text, read_jsonl
 
 
 def json_decode(document, generate):
     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": "'
     s += generate(s, max_tokens=8, stop='"') + '",\n'
     s += '  "country": "'
@@ -24,17 +28,19 @@ def json_decode(document, generate):
     s += generate(s, max_tokens=8, stop='"') + '",\n'
     s += '  "top 3 landmarks": "'
     s += generate(s, max_tokens=24, stop='"') + '",\n'
-    s += '}\n'
+    s += "}\n"
     return s
 
 
 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"],
+            }
+        )
     states = [None] * len(arguments)
 
     # Select backend
@@ -48,13 +54,20 @@ def main(args):
         url = f"{args.host}:{args.port}/generate"
         generate = partial(call_generate_srt_raw, url=url, temperature=0)
     elif args.backend == "guidance":
-        from guidance import models, gen
+        from guidance import gen, models
 
-        model = models.LlamaCpp("/home/ubuntu/model_weights/CodeLlama-7b-instruct-hf.gguf", n_gpu_layers=-1, n_ctx=11000)
+        model = models.LlamaCpp(
+            "/home/ubuntu/model_weights/CodeLlama-7b-instruct-hf.gguf",
+            n_gpu_layers=-1,
+            n_ctx=11000,
+        )
 
         def generate(prompt, max_tokens, stop):
-            out = model + prompt + gen(name="answer",
-                max_tokens=max_tokens, temperature=0, stop=stop)
+            out = (
+                model
+                + prompt
+                + gen(name="answer", max_tokens=max_tokens, temperature=0, stop=stop)
+            )
             return out["answer"]
 
         # warmup
@@ -91,7 +104,7 @@ def main(args):
             "other": {
                 "num_questions": args.num_questions,
                 "parallel": args.parallel,
-            }
+            },
         }
         fout.write(json.dumps(value) + "\n")
 

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")
 

benchmark/long_json_decode/build_dataset.py

@@ -3,7 +3,6 @@ import json
 import transformers
 import wikipedia
 
-
 name = "meta-llama/Llama-2-7b-chat-hf"
 t = transformers.AutoTokenizer.from_pretrained(name)
 city_names = ["los angles", "london", "tokyo", "beijing", "singapore"]
@@ -20,7 +19,9 @@ for city_name in city_names:
     truncate_tokens = t.encode(truncate_content)
 
     # Count token
-    print(f"city_name: {city_name}, #tokens: {len(tokens)}, #truncate tokens: {len(truncate_tokens)}")
+    print(
+        f"city_name: {city_name}, #tokens: {len(tokens)}, #truncate tokens: {len(truncate_tokens)}"
+    )
 
     with open("questions.jsonl", "a") as fout:
         fout.write(json.dumps({"document": truncate_content}) + "\n")