Format Benchmark Code (#399)

2024-04-28 21:06:22 +08:00
parent 19818b9c2f
commit 95c4e0dfac
41 changed files with 1169 additions and 608 deletions
--- a/benchmark/generative_agents/agent_functions.py
+++ b/benchmark/generative_agents/agent_functions.py
@@ -122,16 +122,36 @@ Area options: {Oak Hill College Student Dormatory, The Rose and Crown Pub, Hobbs
 * Must be one of the "Area options," verbatim.
 For eating dinner, Jane Anderson should go to the following area: {Hobbs Cafe}
 ---"""
-    s += (persona_name + " lives in " + living_sector + " that has " +
-          living_sector_areas + ".\n")
-    s += (persona_name + " is currently in " + current_sector + " that has " +
-          current_sector_areas + ".\n")
+    s += (
+        persona_name
+        + " lives in "
+        + living_sector
+        + " that has "
+        + living_sector_areas
+        + ".\n"
+    )
+    s += (
+        persona_name
+        + " is currently in "
+        + current_sector
+        + " that has "
+        + current_sector_areas
+        + ".\n"
+    )
    s += daily_plan + ".\n"
    s += "Area options: " + sector_options + ".\n"
    s += """* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place.
 * Must be one of the "Area options," verbatim.\n"""
-    s += (persona_name + " is " + current_action + ". For " + next_action +
-          ", " + persona_name + " should go to the following area: {")
+    s += (
+        persona_name
+        + " is "
+        + current_action
+        + ". For "
+        + next_action
+        + ", "
+        + persona_name
+        + " should go to the following area: {"
+    )
    s += sgl.gen(name="Location", max_tokens=10, stop="}")


@@ -162,22 +182,43 @@ Area options: {Oak Hill College Student Dormatory, The Rose and Crown Pub, Hobbs
 * Must be one of the "Area options," verbatim.
 For eating dinner, Jane Anderson should go to the following area: {Hobbs Cafe}
 ---"""
-    s += (persona_name + " lives in " + living_sector + " that has " +
-          living_sector_areas + ".\n")
-    s += (persona_name + " is currently in " + current_sector + " that has " +
-          current_sector_areas + ".\n")
+    s += (
+        persona_name
+        + " lives in "
+        + living_sector
+        + " that has "
+        + living_sector_areas
+        + ".\n"
+    )
+    s += (
+        persona_name
+        + " is currently in "
+        + current_sector
+        + " that has "
+        + current_sector_areas
+        + ".\n"
+    )
    s += daily_plan + ".\n"
    s += "Area options: " + sector_options + ".\n"
    s += """* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place.
 * Must be one of the "Area options," verbatim.\n"""
-    s += (persona_name + " is " + current_action + ". For " + next_action +
-          ", " + persona_name + " should go to the following area: {")
+    s += (
+        persona_name
+        + " is "
+        + current_action
+        + ". For "
+        + next_action
+        + ", "
+        + persona_name
+        + " should go to the following area: {"
+    )
    return {"prompt": s, "max_tokens": 10, "stop": "}"}


@sgl.function
-def action_location_object(s, persona_name, target_sector, target_sector_areas,
-                           current_action, next_action):
+def action_location_object(
+    s, persona_name, target_sector, target_sector_areas, current_action, next_action
+):
    s += """
 Jane Anderson is in kitchen in Jane Anderson's house.
 Jane Anderson is going to Jane Anderson's house that has the following areas: {kitchen,  bedroom, bathroom}
@@ -191,20 +232,34 @@ Stay in the current area if the activity can be done there. Never go into other
 For getting coffee, Tom Watson should go to the following area in Hobbs Cafe:
 Answer: {cafe}
 ---"""
-    s += (persona_name + " is going to " + target_sector +
-          " that has the following areas: {" + target_sector_areas + "}\n")
+    s += (
+        persona_name
+        + " is going to "
+        + target_sector
+        + " that has the following areas: {"
+        + target_sector_areas
+        + "}\n"
+    )
    s += """* Stay in the current area if the activity can be done there. 
 * NEVER go into other people's rooms unless necessary."""
-    s += (persona_name + " is " + current_action + ". For " + next_action +
-          ", " + persona_name + "should go to the following area in " +
-          target_sector)
+    s += (
+        persona_name
+        + " is "
+        + current_action
+        + ". For "
+        + next_action
+        + ", "
+        + persona_name
+        + "should go to the following area in "
+        + target_sector
+    )
    s += " (MUST pick one of {" + target_sector_areas + "}):\n"
    s += "Answer: {" + sgl.gen(name="Area", max_tokens=5, stop="}")


-def action_location_object_prompt(persona_name, target_sector,
-                                  target_sector_areas, current_action,
-                                  next_action):
+def action_location_object_prompt(
+    persona_name, target_sector, target_sector_areas, current_action, next_action
+):
    s = ""
    s += """
 Jane Anderson is in kitchen in Jane Anderson's house.
@@ -219,13 +274,27 @@ Stay in the current area if the activity can be done there. Never go into other
 For getting coffee, Tom Watson should go to the following area in Hobbs Cafe:
 Answer: {cafe}
 ---"""
-    s += (persona_name + " is going to " + target_sector +
-          " that has the following areas: {" + target_sector_areas + "}\n")
+    s += (
+        persona_name
+        + " is going to "
+        + target_sector
+        + " that has the following areas: {"
+        + target_sector_areas
+        + "}\n"
+    )
    s += """* Stay in the current area if the activity can be done there. 
 * NEVER go into other people's rooms unless necessary."""
-    s += (persona_name + " is " + current_action + ". For " + next_action +
-          ", " + persona_name + "should go to the following area in " +
-          target_sector)
+    s += (
+        persona_name
+        + " is "
+        + current_action
+        + ". For "
+        + next_action
+        + ", "
+        + persona_name
+        + "should go to the following area in "
+        + target_sector
+    )
    s += " (MUST pick one of {" + target_sector_areas + "}):\n"
    s += "Answer: {"
    return {"prompt": s, "max_tokens": 5, "stop": "}"}
--- a/benchmark/generative_agents/bench_other.py
+++ b/benchmark/generative_agents/bench_other.py
@@ -1,29 +1,29 @@
 import argparse
-from functools import partial
 import json
 import time
+from functools import partial
 from pathlib import Path

+from agent_functions import (
+    action_location_object_prompt,
+    action_location_sector_prompt,
+    generate_event_triple_prompt,
+    generate_pronunciatio_prompt,
+    poignancy_event_prompt,
+)
 from tqdm import tqdm
+
 from sglang.test.test_utils import (
    add_common_other_args_and_parse,
    call_generate_lightllm,
-    call_generate_vllm,
    call_generate_srt_raw,
+    call_generate_vllm,
 )
-from sglang.utils import read_jsonl, dump_state_text
-
-from agent_functions import (
-    poignancy_event_prompt,
-    generate_event_triple_prompt,
-    generate_pronunciatio_prompt,
-    action_location_sector_prompt,
-    action_location_object_prompt,
-)
+from sglang.utils import dump_state_text, read_jsonl


 def main(args):
-    lines = read_jsonl(args.data_path)[:args.num_events]
+    lines = read_jsonl(args.data_path)[: args.num_events]
    mapping = {
        "poignancy_event": poignancy_event_prompt,
        "generate_event_triple": generate_event_triple_prompt,
@@ -46,7 +46,7 @@ def main(args):
        url = f"{args.host}:{args.port}/generate"
        call_generate = partial(call_generate_srt_raw, url=url)
    elif args.backend == "guidance":
-        from guidance import models, gen
+        from guidance import gen, models

        model = models.LlamaCpp(
            str(Path.home()) + "/model_weights/Llama-2-7b-chat.gguf",
@@ -55,11 +55,15 @@ def main(args):
        )

        def call_generate(prompt, temperature, max_tokens, stop):
-            out = model + prompt + gen(
-                name="result",
-                max_tokens=max_tokens,
-                temperature=temperature,
-                stop=stop,
+            out = (
+                model
+                + prompt
+                + gen(
+                    name="result",
+                    max_tokens=max_tokens,
+                    temperature=temperature,
+                    stop=stop,
+                )
            )
            return out["result"]

@@ -87,7 +91,7 @@ def main(args):
            "backend": args.backend,
            "num_gpus": 1,
            "latency": round(latency, 3),
-            # to pack weighted functions as a single agent            
+            # to pack weighted functions as a single agent
            "num_requests": len(arguments) / len(mapping),
            "other": {
                "parallel": args.parallel,
--- a/benchmark/generative_agents/bench_sglang.py
+++ b/benchmark/generative_agents/bench_sglang.py
@@ -2,24 +2,24 @@ import argparse
 import json
 import time

+from agent_functions import (
+    action_location_object,
+    action_location_sector,
+    generate_event_triple,
+    generate_pronunciatio,
+    poignancy_event,
+)
+
 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 agent_functions import (
-    poignancy_event,
-    generate_event_triple,
-    generate_pronunciatio,
-    action_location_sector,
-    action_location_object,
-)
+from sglang.utils import dump_state_text, read_jsonl


 def main(args):
-    lines = read_jsonl(args.data_path)[:args.num_events]
+    lines = read_jsonl(args.data_path)[: args.num_events]
    mapping = {
        "poignancy_event": poignancy_event,
        "generate_event_triple": generate_event_triple,