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