Format Benchmark Code (#399)

benchmark/multi_document_qa/bench_other.py

@@ -1,15 +1,18 @@
 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
 
 USER_PREFIX = "[INST] "
 USER_SUFFIX = " [/INST]"
@@ -25,7 +28,11 @@ def multi_document_qa(docs, question, generate):
     s += "".join(docs)
 
     s += "\nDocuments end."
-    s += ("\n\nBased on the above documents, please answer this question:\n" + question + "\nAnswer in three words or fewer.")
+    s += (
+        "\n\nBased on the above documents, please answer this question:\n"
+        + question
+        + "\nAnswer in three words or fewer."
+    )
     s += USER_SUFFIX
     s += ASSISTANT_PREFIX
     answer = generate(s, max_tokens=16, stop=None)
@@ -42,11 +49,13 @@ def main(args):
     if args.backend == "guidance":
         num_docs = 7  # due to OOM
 
-    for i in range(len(l["questions"][:args.num_questions])):
-        arguments.append({
-            "docs": l["documents"][:num_docs],
-            "question": l["questions"][i],
-        })
+    for i in range(len(l["questions"][: args.num_questions])):
+        arguments.append(
+            {
+                "docs": l["documents"][:num_docs],
+                "question": l["questions"][i],
+            }
+        )
         labels.append(l["answers"][i])
     states = [None] * len(arguments)
 
@@ -61,13 +70,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
@@ -113,7 +129,7 @@ def main(args):
             "other": {
                 "num_questions": args.num_questions,
                 "parallel": args.parallel,
-            }
+            },
         }
         fout.write(json.dumps(value) + "\n")