Add city doc benchmark mode (#129)

benchmark/json_fast_forward/bench_other.py

@@ -9,7 +9,7 @@ from sglang.test.test_utils import (
     add_common_other_args_and_parse,
     call_generate_outlines,
 )
-from sglang.utils import dump_state_text
+from sglang.utils import dump_state_text, read_jsonl
 from tqdm import tqdm
 
 # there are some FSM bugs with json regex converted from pydantic model
@@ -32,6 +32,16 @@ character_regex = (
     + r"""\}"""
 )
 
+city_regex = (
+    r"""\{\n"""
+    + r"""  "name": "[\w\d\s]{1,16}",\n"""
+    + r"""  "country": "[\w\d\s]{1,16}",\n"""
+    + r"""  "latitude": [-+]?[0-9]*\.?[0-9]{0,2},\n"""
+    + r"""  "population": [-+]?[0-9]{1,9},\n"""
+    + r"""  "top 3 landmarks": \["[\w\d\s]{1,16}", "[\w\d\s]{1,16}", "[\w\d\s]{1,16}"\]\n"""
+    + r"""\}"""
+)
+
 # fmt: off
 def character_gen(name, generate):
     s = name + " is a character in Harry Potter. Please fill in the following information about this character.\n"
@@ -39,6 +49,15 @@ def character_gen(name, generate):
     return s
 # fmt: on
 
+# fmt: off
+def city_gen(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 += generate(s, max_tokens=256, regex=city_regex)
+    return s
+# fmt: on
+
 
 @guidance
 def character_maker(lm, name):
@@ -65,7 +84,31 @@ def character_maker(lm, name):
     return lm
 
 
-def main(args):
+@guidance
+def city_maker(lm, document):
+    regex_str_no_quote = r"[\w\d\s]+"
+    regex_float = r"[0-9]+\.[0-9]+"
+    lm += f"""\
+    Please extract the information of a city from the following wikipedia page.
+    Page begin.
+    {document}
+    Page end.
+    Here is the name, country, and symbol of the city in JSON format.
+    {{
+        "name": "{guidance.gen("name", max_tokens=16, regex=regex_str_no_quote)}",
+        "country": "{guidance.gen("country", max_tokens=16, regex=regex_str_no_quote)}",
+        "latitude": {guidance.gen("latitude", max_tokens=10, regex=regex_float)},
+        "population": {guidance.gen("population", max_tokens=10, regex=r"[0-9]+")},
+        "top 3 landmarks": [
+            "{guidance.gen("landmark1", max_tokens=16, regex=regex_str_no_quote)}", "{guidance.gen("landmark2", max_tokens=16, regex=regex_str_no_quote)}", "{guidance.gen("landmark3", max_tokens=16, regex=regex_str_no_quote)}"
+        ]
+    }}
+    """
+
+    return lm
+
+
+def bench_character(args):
     arguments = []
     with open(args.data_path, "r") as f:
         for line in f:
@@ -85,7 +128,7 @@ def main(args):
         get_one_answer = func
     elif args.backend == "guidance":
         model = guidance.models.LlamaCpp(
-            "/home/ubuntu/model_weights/Llama-2-7b-chat-hf/ggml-model-f16.gguf",
+            args.llama_cpp_model_path,
             n_gpu_layers=-1,
             n_ctx=4096,
         )
@@ -110,11 +153,69 @@ def main(args):
 
     latency = time.time() - tic
 
+    return states, latency
+
+
+def bench_city_doc(args):
+    arguments = []
+    for line in read_jsonl(args.data_path):
+        arguments.append({"document": line["document"]})
+    arguments = arguments[: args.num_jsons]
+
+    states = [None] * len(arguments)
+
+    # Select backend
+    if args.backend == "vllm":
+        url = f"{args.host}:{args.port}/generate"
+        generate = partial(call_generate_outlines, url=url, temperature=0)
+
+        def func(i):
+            states[i] = city_gen(**arguments[i], generate=generate)
+
+        get_one_answer = func
+    elif args.backend == "guidance":
+        model = guidance.models.LlamaCpp(
+            args.llama_cpp_model_path,
+            n_gpu_layers=-1,
+            n_ctx=4096,
+        )
+
+        def func(i):
+            lm = model + city_maker(**arguments[i])
+            states[i] = lm
+
+        get_one_answer = func
+    else:
+        raise ValueError(f"Invalid backend: {args.backend}")
+
+    tic = time.time()
+    if args.parallel == 1:
+        for i in tqdm(range(len(arguments))):
+            get_one_answer(i)
+    else:
+        with ThreadPoolExecutor(args.parallel) as executor:
+            rets = executor.map(get_one_answer, list(range(len(arguments))))
+            for _ in rets:
+                pass
+
+    latency = time.time() - tic
+
+    return states, latency
+
+
+def main(args):
+    if args.mode == "character":
+        args.data_path = "dataset.txt"
+        states, latency = bench_character(args)
+    elif args.mode == "city":
+        args.data_path = "questions.jsonl"
+        states, latency = bench_city_doc(args)
+
     # Compute accuracy
     print(f"Latency: {latency:.3f}")
 
     # Write results
-    dump_state_text(f"tmp_output_{args.backend}.txt", states)
+    dump_state_text(f"tmp_output_{args.backend}_{args.mode}.txt", states)
 
     with open(args.result_file, "a") as fout:
         value = {
@@ -129,7 +230,15 @@ def main(args):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--data-path", type=str, default="dataset.txt")
+    parser.add_argument("--data-path", type=str)
     parser.add_argument("--num-jsons", type=int, default=50)
+    parser.add_argument(
+        "--mode", type=str, default="character", choices=["character", "city"]
+    )
+    parser.add_argument(
+        "--llama-cpp-model-path",
+        type=str,
+        default="/home/ubuntu/model_weights/Llama-2-7b-chat-hf/ggml-model-f16.gguf",
+    )
     args = add_common_other_args_and_parse(parser)
     main(args)