release initial code

Co-authored-by: Ying Sheng <sqy1415@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
Co-authored-by: parasol-aser <3848358+parasol-aser@users.noreply.github.com>
Co-authored-by: LiviaSun <33578456+ChuyueSun@users.noreply.github.com>
Co-authored-by: Cody Yu <hao.yu.cody@gmail.com>

benchmark/multi_chain_reasoning/README.md

+## Download data
+```
+wget https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl
+```
+
+## Run benchmark
+
+### Benchmark sglang
+```
+python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000
+```
+
+```
+python3 bench_sglang.py --num-questions 64
+python3 bench_sglang.py --num-questions 32 --parallel 1
+```
+
+
+### Benchmark vllm
+```
+python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000
+```
+
+```
+python3 bench_other.py --num-questions 64 --backend vllm
+```
+
+
+### Benchmark lightllm
+```
+# A10G
+python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000
+```
+
+```
+python3 bench_other.py --num-questions 64 --backend lightllm
+```
+
+
+### Benchmark guidance
+```
+python3 bench_other.py --num-questions 8 --backend guidance --parallel 1
+```

benchmark/multi_chain_reasoning/bench_other.py

+import argparse
+import ast
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+import json
+import re
+import time
+
+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
+
+
+INVALID = -9999999
+
+
+def get_answer_value(answer_str):
+    answer_str = answer_str.replace(",", "")
+    numbers = re.findall(r'\d+', answer_str)
+    if len(numbers) < 1:
+        return INVALID
+    try:
+        return ast.literal_eval(numbers[-1])
+    except SyntaxError:
+        return INVALID
+
+
+prompt_lib = [
+    "Let us think step by step.",
+    "Approach this methodically. Let's dissect the problem into smaller, more manageable parts.",
+    "It's important to proceed step by step, ensuring accuracy at each stage.",
+    "Take a deep breath and break this down.",
+    "A little bit of arithmetic and a logical approach will help us quickly arrive at the solution to this problem.",
+    "I am extremely good at math.",
+]
+
+
+def multi_chain_gsm8k(question, num_chains, call_generate):
+    s = "Question: " + question + "\n"
+    # s += call_generate(s + "Answer: " + prompt_lib[0], max_tokens=256,
+    #     stop="Question", temperature=0)
+    # return s
+
+    comps = []
+    for i in range(num_chains):
+        comps.append(call_generate(s + "Answer: " + prompt_lib[i % num_chains],
+                     max_tokens=256, temperature=0.3, stop="Question"))
+
+    s += "Answer: To answer this question, here are some possible solutions. "
+    s += "After considering all of them, I will do a majority vote.\n\n"
+    for i in range(num_chains):
+        s += f"Solution {i+1}: " + comps[i].strip() + "\n\n"
+    s += f"\nBy considering the above solutions and doing a majority vote, I think the final answer (a single integer number) is "
+    s += call_generate(s, max_tokens=16, temperature=0, stop=None)
+    return s
+
+
+def main(args):
+    lines = read_jsonl(args.data_path)
+
+    # Construct prompts
+    k = args.num_shot
+
+    questions = []
+    labels = []
+    for i in range(len(lines[:args.num_questions])):
+        questions.append(lines[i]["question"])
+        labels.append(get_answer_value(lines[i]["answer"]))
+    assert all(l != INVALID for l in labels)
+
+    states = [None] * len(labels)
+
+    # Select backend
+    if args.backend == "lightllm":
+        url = f"{args.host}:{args.port}/generate"
+        call_generate = partial(call_generate_lightllm, url=url)
+    elif args.backend == "vllm":
+        url = f"{args.host}:{args.port}/generate"
+        call_generate = partial(call_generate_vllm, url=url)
+    elif args.backend == "srt-raw":
+        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
+
+        model = models.LlamaCpp("/home/ubuntu/model_weights/Llama-2-7b-chat.gguf", n_gpu_layers=-1, n_ctx=4096)
+
+        def call_generate(prompt, temperature, max_tokens, stop):
+            out = model + prompt + gen(name="answer",
+                max_tokens=max_tokens, temperature=temperature, stop=stop)
+            return out["answer"]
+
+        #def multi_chain_gsm8k(question, num_chains, call_generate):
+        #    s = model + "Question: " + question + "\n"
+
+        #    comps = []
+        #    for i in range(num_chains):
+        #        comps.append(call_generate(s + "Answer: " + prompt_lib[i % num_chains],
+        #                     max_tokens=256, temperature=0.3, stop="Question"))
+
+        #    s += "Answer: To answer this question, here are some possible solutions. "
+        #    s += "After considering all of them, I will do a majority vote.\n\n"
+        #    for i in range(num_chains):
+        #        s += f"Solution {i+1}: " + comps[i].strip() + "\n\n"
+        #    s += f"\nBy considering the above solutions and doing a majority vote, I think the final answer (a single integer number) is "
+        #    return call_generate(s, max_tokens=16, temperature=0, stop=None)
+
+    elif args.backend == "lmql":
+        import lmql
+        model = lmql.model("meta-llama/Llama-2-7b-chat-hf",
+           endpoint=f"{args.host}:{args.port}")
+
+        @lmql.query(model=model)
+        async def program(question):
+            '''lmql
+            """{question}[ANSWER]""" where len(TOKENS(ANSWER)) < 257 and STOPS_AT(ANSWER, "Question")
+            return ANSWER
+            '''
+
+        async def call_generate(prompt, temperature, max_tokens, stop):
+            return await program(question=prompt, temperature=0)
+
+    else:
+        raise ValueError(f"Invalid backend: {args.backend}")
+
+    # Run requests
+    if args.backend != "lmql":
+        # Use thread pool
+        def get_one_answer(i):
+            answer = multi_chain_gsm8k(questions[i], args.num_chains,
+                call_generate)
+            states[i] = answer
+
+        tic = time.time()
+        if args.parallel == 1:
+            for i in range(len(questions)):
+                get_one_answer(i)
+        else:
+            with ThreadPoolExecutor(args.parallel) as executor:
+                executor.map(get_one_answer, list(range(len(questions))))
+    else:
+        # Use asyncio
+        async def batched_call(batch_size):
+            for i in range(0, len(questions), batch_size):
+                tasks = []
+                for q in questions[i:i+batch_size]:
+                    tasks.append(call_generate(few_shot_examples + q,
+                        temperature=0, max_tokens=256, stop="Question"))
+                rets = await asyncio.gather(*tasks)
+                for j in range(len(rets)):
+                    states[i+j] = get_answer_value(rets[j])
+
+        tic = time.time()
+        asyncio.run(batched_call(batch_size=args.parallel))
+    latency = time.time() - tic
+
+    preds = []
+    for i in range(len(states)):
+        preds.append(get_answer_value(states[i]))
+
+    # Compute accuracy
+    acc = np.mean(np.array(preds) == np.array(labels))
+    invalid = np.mean(np.array(preds) == INVALID)
+    print(f"Latency: {latency:.3f}")
+    print(f"Invalid: {invalid:.3f}")
+    print(f"Accuracy: {acc:.3f}")
+
+    # Write results
+    dump_state_text(f"tmp_output_{args.backend}.txt", states)
+
+    with open(args.result_file, "a") as fout:
+        value = {
+            "task": "multi_chain_gsm8k",
+            "backend": args.backend,
+            "num_gpus": 1,
+            "latency": round(latency, 3),
+            "accuracy": round(acc, 3),
+            "num_requests": args.num_questions,
+            "other": {
+                "num_questions": args.num_questions,
+                "parallel": args.parallel,
+            }
+        }
+        fout.write(json.dumps(value) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num-shot", type=int, default=0)
+    parser.add_argument("--num-chains", type=int, default=5)
+    parser.add_argument("--data-path", type=str, default="test.jsonl")
+    parser.add_argument("--num-questions", type=int, default=50)
+    args = add_common_other_args_and_parse(parser)
+    main(args)

benchmark/multi_chain_reasoning/bench_sglang.py

+import argparse
+import ast
+import json
+import re
+import time
+
+import numpy as np
+from sglang.test.test_utils import add_common_sglang_args_and_parse, select_sglang_backend
+from sglang.utils import read_jsonl, dump_state_text
+
+
+INVALID = -9999999
+
+
+def get_answer_value(answer_str):
+    answer_str = answer_str.replace(",", "")
+    numbers = re.findall(r'\d+', answer_str)
+    if len(numbers) < 1:
+        return INVALID
+    try:
+        return ast.literal_eval(numbers[-1])
+    except SyntaxError:
+        return INVALID
+
+
+prompt_lib = [
+    "Let us think step by step.",
+    "Approach this methodically. Let's dissect the problem into smaller, more manageable parts.",
+    "It's important to proceed step by step, ensuring accuracy at each stage.",
+    "Take a deep breath and break this down.",
+    "A little bit of arithmetic and a logical approach will help us quickly arrive at the solution to this problem.",
+    "I am extremely good at math.",
+]
+
+
+def main(args):
+    lines = read_jsonl(args.data_path)
+
+    # Construct prompts
+    #k = args.num_shot
+    #few_shot_examples = get_few_shot_examples(lines, k)
+
+    questions = []
+    labels = []
+    for i in range(len(lines[:args.num_questions])):
+        questions.append(lines[i]["question"])
+        labels.append(get_answer_value(lines[i]["answer"]))
+    assert all(l != INVALID for l in labels)
+    arguments = [{"question": q} for q in questions]
+
+    num_chains = args.num_chains
+
+    #####################################
+    ######### SGL Program Begin #########
+    #####################################
+
+    import sglang as sgl
+
+    @sgl.function
+    def multi_chain_gsm8k(s, question):
+        s += "Question: " + question + "\n"
+        #s += "Answer: " + prompt_lib[0] + sgl.gen("answer", max_tokens=256, stop="Question",
+        #    temperature=0)
+        #return
+
+        forks = s.fork(num_chains)
+        for i in range(num_chains):
+            forks[i] += ("Answer: " + prompt_lib[i % num_chains] +
+                sgl.gen(f"chain", max_tokens=256, temperature=0.3, stop="Question"))
+        forks.join()
+
+        s += "Answer: To answer this question, here are some possible solutions. "
+        s += "After considering all of them, I will do a majority vote.\n\n"
+        for i in range(num_chains):
+            s += f"Solution {i+1}: " + forks[i]["chain"].strip() + "\n\n"
+        s += f"\nBy considering the above solutions and doing a majority vote, I think the final answer (a single integer number) is "
+        s += sgl.gen("answer", max_tokens=16)
+
+    #####################################
+    ########## SGL Program End ##########
+    #####################################
+
+    # Select backend
+    backend = select_sglang_backend(args)
+
+    # Run requests
+    tic = time.time()
+    states = multi_chain_gsm8k.run_batch(
+        arguments, temperature=0, backend=backend, num_threads=args.parallel)
+    latency = time.time() - tic
+
+    preds = []
+    for i in range(len(states)):
+        preds.append(get_answer_value(states[i]["answer"]))
+
+    # Compute accuracy
+    acc = np.mean(np.array(preds) == np.array(labels))
+    invalid = np.mean(np.array(preds) == INVALID)
+    print(f"Latency: {latency:.3f}")
+    print(f"Invalid: {invalid:.3f}")
+    print(f"Accuracy: {acc:.3f}")
+
+    # Write results
+    dump_state_text(f"tmp_output_{args.backend}.txt", states)
+
+    with open(args.result_file, "a") as fout:
+        value = {
+            "task": "multi_chain_gsm8k",
+            "backend": args.backend,
+            "num_gpus": 1,
+            "latency": round(latency, 3),
+            "accuracy": round(acc, 3),
+            "num_requests": args.num_questions,
+            "other": {
+                "num_questions": args.num_questions,
+                "parallel": args.parallel,
+            }
+        }
+        fout.write(json.dumps(value) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num-shot", type=int, default=0)
+    parser.add_argument("--num-chains", type=int, default=5)
+    parser.add_argument("--data-path", type=str, default="test.jsonl")
+    parser.add_argument("--num-questions", type=int, default=50)
+    args = add_common_sglang_args_and_parse(parser)
+    main(args)

benchmark/multi_document_qa/README.md

+## Run benchmark
+
+### Benchmark sglang
+```
+python3 -m sglang.launch_server --model-path codellama/CodeLlama-7b-instruct-hf --port 30000
+```
+
+```
+python3 bench_sglang.py --num-questions 10 --parallel 1
+```
+
+
+### Benchmark vllm
+```
+python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model codellama/CodeLlama-7b-instruct-hf  --disable-log-requests --port 21000 --gpu 0.97
+```
+
+```
+python3 bench_other.py --backend vllm --num-questions 64
+```
+
+
+### Benchmark guidance
+```
+python3 bench_other.py --backend guidance --num-questions 32 --parallel 1
+```
+
+
+
+### Build dataset
+
+```
+pip install PyPDF2
+python3 build_dataset.py
+```
+
+```python
+import PyPDF2
+
+with open('llama2.pdf', 'rb') as file:
+    reader = PyPDF2.PdfReader(file)
+    text = ''
+    for page_num in range(len(reader.pages)):
+        text += reader.pages[page_num].extract_text()
+    with open('output.txt', 'w') as text_file:
+        text_file.write(text)
+```

benchmark/multi_document_qa/bench_other.py

+import argparse
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+import json
+import time
+
+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
+
+
+USER_PREFIX = "[INST] "
+USER_SUFFIX = " [/INST]"
+ASSISTANT_PREFIX = ""
+ASSISTANT_SUFFIX = " </s><s>"
+
+
+def multi_document_qa(docs, question, generate):
+    s = USER_PREFIX
+    s += "Pleaes answer a question according to given documents.\n"
+    s += "Question:" + question + "Documents begin.\n"
+
+    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 += USER_SUFFIX
+    s += ASSISTANT_PREFIX
+    answer = generate(s, max_tokens=16, stop=None)
+    return answer
+
+
+def main(args):
+    lines = read_jsonl(args.data_path)
+    l = lines[0]
+    arguments = []
+    labels = []
+
+    num_docs = 10
+    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],
+        })
+        labels.append(l["answers"][i])
+    states = [None] * len(arguments)
+
+    # Select backend
+    if args.backend == "lightllm":
+        url = f"{args.host}:{args.port}/generate"
+        generate = partial(call_generate_lightllm, url=url, temperature=0)
+    elif args.backend == "vllm":
+        url = f"{args.host}:{args.port}/generate"
+        generate = partial(call_generate_vllm, url=url, temperature=0)
+    elif args.backend == "srt-raw":
+        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
+
+        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)
+            return out["answer"]
+
+        # warmup
+        generate("Hello!", max_tokens=8, stop=None)
+    else:
+        raise ValueError(f"Invalid backend: {args.backend}")
+
+    # Run requests
+    def get_one_answer(i):
+        states[i] = multi_document_qa(generate=generate, **arguments[i])
+
+    tic = time.time()
+    if args.parallel == 1:
+        for i in tqdm(range(len(labels))):
+            get_one_answer(i)
+    else:
+        with ThreadPoolExecutor(args.parallel) as executor:
+            executor.map(get_one_answer, list(range(len(labels))))
+    latency = time.time() - tic
+
+    # Compute accuracy
+    print(states)
+    correct = 0
+    for s, label in zip(states, labels):
+        answer = s.lower()
+        if all(x in answer for x in label.lower().split(" ")):
+            correct += 1
+    accuracy = correct / len(labels)
+    print(f"Accuracy: {accuracy:.3f}")
+    print(f"Latency: {latency:.3f}")
+
+    # Write results
+    dump_state_text(f"tmp_output_{args.backend}.txt", states)
+
+    with open(args.result_file, "a") as fout:
+        value = {
+            "task": "multi_document_qa",
+            "backend": args.backend,
+            "num_gpus": 1,
+            "latency": round(latency, 3),
+            "num_requests": args.num_questions,
+            "accuracy": accuracy,
+            "other": {
+                "num_questions": args.num_questions,
+                "parallel": args.parallel,
+            }
+        }
+        fout.write(json.dumps(value) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-path", type=str, default="questions.jsonl")
+    parser.add_argument("--num-questions", type=int, default=100)
+    args = add_common_other_args_and_parse(parser)
+    main(args)

benchmark/multi_document_qa/bench_sglang.py

+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
+
+
+@sgl.function
+def multi_document_qa(s, docs, question):
+    s += sgl.user_begin()
+    s += "Pleaes answer a question according to given documents.\n"
+    s += "Question:" + question + "Documents begin.\n"
+
+    forks = s.fork(len(docs))
+    forks += lambda i: docs[i]
+    forks.join("concate_and_append")
+
+    s += "\nDocuments end."
+    s += ("\n\nBased on the above documents, please answer this question:\n" + question + "\nAnswer in three words or fewer.")
+    s += sgl.user_end()
+    s += sgl.assistant(sgl.gen("answer", max_tokens=16))
+
+
+def main(args):
+    lines = read_jsonl(args.data_path)
+    l = lines[0]
+    arguments = []
+    labels = []
+    for i in range(len(l["questions"][:args.num_questions])):
+        arguments.append({
+            "docs": l["documents"][:10],
+            "question": l["questions"][i],
+        })
+        labels.append(l["answers"][i])
+
+    # Select backend
+    backend = select_sglang_backend(args)
+    sgl.set_default_backend(backend)
+
+    # Run requests
+    tic = time.time()
+    states = multi_document_qa.run_batch(
+        arguments, temperature=0, num_threads=args.parallel)
+    latency = time.time() - tic
+
+    # Compute accuracy 
+    print([s["answer"] for s in states])
+    correct = 0
+    for s, label in zip(states, labels):
+        answer = s["answer"].lower()
+        if all(x in answer for x in label.lower().split(" ")):
+            correct += 1
+    accuracy = correct / len(labels)
+    print(f"Accuracy: {accuracy:.3f}")
+    print(f"Latency: {latency:.3f}")
+
+    # Write results
+    dump_state_text(f"tmp_output_{args.backend}.txt", states)
+
+    with open(args.result_file, "a") as fout:
+        value = {
+            "task": "multi_document_qa",
+            "backend": args.backend,
+            "num_gpus": 1,
+            "latency": round(latency, 3),
+            "num_requests": args.num_questions,
+            "accuracy": accuracy,
+            "other": {
+                "num_questions": args.num_questions,
+                "parallel": args.parallel,
+            }
+        }
+        fout.write(json.dumps(value) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-path", type=str, default="questions.jsonl")
+    parser.add_argument("--num-questions", type=int, default=100)
+    args = add_common_sglang_args_and_parse(parser)
+    main(args)

benchmark/multi_document_qa/build_dataset.py

+import json
+
+import transformers
+
+content = "\n".join(
+    open("llama2.txt", 'r', encoding='utf-8', errors='ignore').readlines())
+content = content.replace("\n\n", "\n")
+
+# Count token
+name = "meta-llama/Llama-2-7b-chat-hf"
+t = transformers.AutoTokenizer.from_pretrained(name)
+print(f"num tokens: {len(t.encode(content))}")
+
+# Segment
+SEP = "\n\n"
+parts = content.split(SEP)
+print(f"num segments: {len(parts)}")
+
+segment_len = 1100
+
+segments = []
+tmp = []
+tmp_len = 0
+for i in range(len(parts)):
+    tmp.append(parts[i])
+    tmp_len += len(t.encode(parts[i]))
+
+    if tmp_len > segment_len:
+        segments.append(SEP.join(tmp))
+        tmp = []
+        tmp_len = 0
+
+for i, s in enumerate(segments):
+    print(i, len(t.encode(segments[i])))
+
+# Dump
+with open("questions.jsonl", "w") as fout:
+    fout.write(json.dumps({
+        "documents": segments[:30],
+        "questions": [
+            "What is the name of the fine-tuned LLMs?",
+            "Which figure shows the helpfulness human evaluation results for Llama 2-Chat?",
+            "What is the number of parameters in the largest Llama 2 model?",
+            "What is the batch size of fine-tuning?",
+            "Where can we find the details of potential data contamination?",
+            "What is the full name of MPT?",
+            "What is the power consumption of RSC in Watt?",
+            "How many tokens of data do they train on?",
+            "Which model's release is delayed due to a lack of time to sufficiently red team?",
+            "Which activation function is used in Llama?"
+        ],
+        "answers": [
+            "Llama 2 Chat",
+            "1",
+            "70 B",
+            "64",
+            "A 6",
+            "MosaicML",
+            "400",
+            "2 trillion",
+            "34 B",
+            "SwiGLU",
+        ],
+    }) + "\n")

benchmark/react/README.md

+## Run benchmark
+
+### Benchmark sglang
+```
+python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000
+```
+
+```
+python3 bench_sglang.py --num-questions 100
+```
+
+
+### Benchmark vllm
+```
+python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000
+```
+
+```
+python3 bench_other.py --num-questions 100 --backend vllm
+```
+
+
+### Benchmark guidance
+```
+python3 bench_other.py --num-questions 100 --backend guidance --parallel 1
+```

benchmark/react/bench_other.py

+import argparse
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+import json
+import time
+from pathlib import Path
+
+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,
+)
+from sglang.utils import read_jsonl, dump_state_text
+
+
+def get_prompt(question):
+    prompt = (
+        """Solve a question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types: 
+(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search.
+(2) Lookup[keyword], which returns the next sentence containing keyword in the current passage.
+(3) Finish[answer], which returns the answer and finishes the task.
+Here are some examples.
+Question: What is the elevation range for the area that the eastern sector of the Colorado orogeny extends into?
+Thought 1: I need to search Colorado orogeny, find the area that the eastern sector of the Colorado orogeny extends into, then find the elevation range of the area.
+Action 1: Search[Colorado orogeny]
+Observation 1: The Colorado orogeny was an episode of mountain building (an orogeny) in Colorado and surrounding areas.
+Thought 2: It does not mention the eastern sector. So I need to look up eastern sector.
+Action 2: Lookup[eastern sector]
+Observation 2: (Result 1 / 1) The eastern sector extends into the High Plains and is called the Central Plains orogeny.
+Thought 3: The eastern sector of Colorado orogeny extends into the High Plains. So I need to search High Plains and find its elevation range.
+Action 3: Search[High Plains]
+Observation 3: High Plains refers to one of two distinct land regions:
+Thought 4: I need to instead search High Plains (United States).
+Action 4: Search[High Plains (United States)]
+Observation 4: The High Plains are a subregion of the Great Plains. From east to west, the High Plains rise in elevation from around 1,800 to 7,000 ft (550 to 2,130 m).[3]
+Thought 5: High Plains rise in elevation from around 1,800 to 7,000 ft, so the answer is 1,800 to 7,000 ft.
+Action 5: Finish[1,800 to 7,000 ft]
+Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who?
+Thought 1: The question simplifies to "The Simpsons" character Milhouse is named after who. I only need to search Milhouse and find who it is named after.
+Action 1: Search[Milhouse]
+Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening.
+Thought 2: The paragraph does not tell who Milhouse is named after, maybe I can look up "named after".
+Action 2: Lookup[named after]
+Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous. 
+Thought 3: Milhouse was named after U.S. president Richard Nixon, so the answer is Richard Nixon.
+Action 3: Finish[Richard Nixon]
+Question: Which documentary is about Finnish rock groups, Adam Clayton Powell or The Saimaa Gesture?
+Thought 1: I need to search Adam Clayton Powell and The Saimaa Gesture, and find which documentary is about Finnish rock groups.
+Action 1: Search[Adam Clayton Powell]
+Observation 1: Could not find [Adam Clayton Powell]. Similar: ['Adam Clayton Powell III', 'Seventh Avenue (Manhattan)', 'Adam Clayton Powell Jr. State Office Building', 'Isabel Washington Powell', 'Adam Powell', 'Adam Clayton Powell (film)', 'Giancarlo Esposito'].
+Thought 2: To find the documentary, I can search Adam Clayton Powell (film).
+Action 2: Search[Adam Clayton Powell (film)]
+Observation 2: Adam Clayton Powell is a 1989 American documentary film directed by Richard Kilberg.
+The film is about the rise and fall of influential African-American politician Adam Clayton Powell Jr.[3][4] It was later aired as part of the PBS series The American Experience.
+Thought 3: Adam Clayton Powell (film) is a documentary about an African-American politician, not Finnish rock groups. So the documentary about Finnish rock groups must instead be The Saimaa Gesture.
+Action 3: Finish[The Saimaa Gesture]
+Question: What profession does Nicholas Ray and Elia Kazan have in common?
+Thought 1: I need to search Nicholas Ray and Elia Kazan, find their professions, then find the profession they have in common.
+Action 1: Search[Nicholas Ray]
+Observation 1: Nicholas Ray (born Raymond Nicholas Kienzle Jr., August 7, 1911 – June 16, 1979) was an American film director, screenwriter, and actor best known for the 1955 film Rebel Without a Cause.
+Thought 2: Professions of Nicholas Ray are director, screenwriter, and actor. I need to search Elia Kazan next and find his professions.
+Action 2: Search[Elia Kazan]
+Observation 2: Elia Kazan was an American film and theatre director, producer, screenwriter and actor.
+Thought 3: Professions of Elia Kazan are director, producer, screenwriter, and actor. So profession Nicholas Ray and Elia Kazan have in common is director, screenwriter, and actor.
+Action 3: Finish[director, screenwriter, actor]
+Question: Which magazine was started first Arthur's Magazine or First for Women?
+Thought 1: I need to search Arthur's Magazine and First for Women, and find which was started first.
+Action 1: Search[Arthur's Magazine]
+Observation 1: Arthur's Magazine (1844-1846) was an American literary periodical published in Philadelphia in the 19th century. 
+Thought 2: Arthur's Magazine was started in 1844. I need to search First for Women next.
+Action 2: Search[First for Women]
+Observation 2: First for Women is a woman's magazine published by Bauer Media Group in the USA.[1] The magazine was started in 1989. 
+Thought 3: First for Women was started in 1989. 1844 (Arthur's Magazine) < 1989 (First for Women), so Arthur's Magazine was started first.
+Action 3: Finish[Arthur's Magazine]
+Question: Were Pavel Urysohn and Leonid Levin known for the same type of work?
+Thought 1: I need to search Pavel Urysohn and Leonid Levin, find their types of work, then find if they are the same.
+Action 1: Search[Pavel Urysohn]
+Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory.
+Thought 2: Pavel Urysohn is a mathematician. I need to search Leonid Levin next and find its type of work.
+Action 2: Search[Leonid Levin]
+Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist. 
+Thought 3: Leonid Levin is a mathematician and computer scientist. So Pavel Urysohn and Leonid Levin have the same type of work. 
+Action 3: Finish[yes]
+""" + question)
+    return prompt
+
+
+def main(args):
+    lines = read_jsonl(args.data_path)[:args.num_questions]
+    arguments = [{
+        "question": k,
+        "triplets": v
+    } for l in lines for k, v in l.items()]
+
+    states = []
+
+    # Select backend
+    if args.backend == "lightllm":
+        url = f"{args.host}:{args.port}/generate"
+        call_generate = partial(call_generate_lightllm, url=url)
+    elif args.backend == "vllm":
+        url = f"{args.host}:{args.port}/generate"
+        call_generate = partial(call_generate_vllm, url=url)
+    elif args.backend == "srt-raw":
+        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
+
+        model = models.LlamaCpp(
+            str(Path.home()) + "/model_weights/Llama-2-7b-chat.gguf",
+            n_gpu_layers=-1,
+            n_ctx=4096,
+        )
+
+        def call_generate(prompt, temperature, max_tokens, stop):
+            out = (model + prompt + gen(
+                name="result",
+                max_tokens=max_tokens,
+                temperature=temperature,
+                stop=stop,
+            ))
+            return out["result"]
+
+    else:
+        raise ValueError(f"Invalid backend: {args.backend}")
+
+    def run_single_agent(argument):
+        question = argument["question"]
+        triplets = argument["triplets"]
+        prompt = get_prompt(question)
+        for i in range(1, len(triplets) + 2):
+            prompt += "Thought " + str(i) + ":"
+            states.append(prompt)
+            answer = call_generate(prompt,
+                                   max_tokens=200,
+                                   temperature=0,
+                                   stop="Observation")
+            if i > len(triplets):
+                break
+            prompt += (triplets[i - 1]["thought"] + "\nAction " + str(i) +
+                       ":" + triplets[i - 1]["action"] + "\nObservation " +
+                       str(i) + ":" + triplets[i - 1]["observation"] + "\n")
+
+            states.append(answer)
+
+    tic = time.time()
+    if args.parallel == 1:
+        for arg in tqdm(arguments):
+            run_single_agent(arg)
+    else:
+        with ThreadPoolExecutor(args.parallel) as executor:
+            executor.map(run_single_agent, arguments)
+    latency = time.time() - tic
+
+    print(f"Latency: {latency:.3f}")
+
+    # Write results
+    dump_state_text(f"tmp_output_{args.backend}.txt", states)
+
+    with open(args.result_file, "a") as fout:
+        value = {
+            "task": "ReAct Agents",
+            "backend": args.backend,
+            "num_gpus": 1,
+            "latency": round(latency, 3),
+            "num_requests": len(arguments),
+            "other": {
+                "parallel": args.parallel,
+            },
+        }
+        fout.write(json.dumps(value) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-path", type=str, default="hotpotqa_100.jsonl")
+    parser.add_argument("--num-questions", type=int, default=10)
+    args = add_common_other_args_and_parse(parser)
+    main(args)

benchmark/react/bench_sglang.py

+import argparse
+import json
+import time
+
+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
+
+
+@sgl.function
+def webthink(s, question, triplets):
+    s += (
+        """Solve a question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types: 
+(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search.
+(2) Lookup[keyword], which returns the next sentence containing keyword in the current passage.
+(3) Finish[answer], which returns the answer and finishes the task.
+Here are some examples.
+Question: What is the elevation range for the area that the eastern sector of the Colorado orogeny extends into?
+Thought 1: I need to search Colorado orogeny, find the area that the eastern sector of the Colorado orogeny extends into, then find the elevation range of the area.
+Action 1: Search[Colorado orogeny]
+Observation 1: The Colorado orogeny was an episode of mountain building (an orogeny) in Colorado and surrounding areas.
+Thought 2: It does not mention the eastern sector. So I need to look up eastern sector.
+Action 2: Lookup[eastern sector]
+Observation 2: (Result 1 / 1) The eastern sector extends into the High Plains and is called the Central Plains orogeny.
+Thought 3: The eastern sector of Colorado orogeny extends into the High Plains. So I need to search High Plains and find its elevation range.
+Action 3: Search[High Plains]
+Observation 3: High Plains refers to one of two distinct land regions:
+Thought 4: I need to instead search High Plains (United States).
+Action 4: Search[High Plains (United States)]
+Observation 4: The High Plains are a subregion of the Great Plains. From east to west, the High Plains rise in elevation from around 1,800 to 7,000 ft (550 to 2,130 m).[3]
+Thought 5: High Plains rise in elevation from around 1,800 to 7,000 ft, so the answer is 1,800 to 7,000 ft.
+Action 5: Finish[1,800 to 7,000 ft]
+Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who?
+Thought 1: The question simplifies to "The Simpsons" character Milhouse is named after who. I only need to search Milhouse and find who it is named after.
+Action 1: Search[Milhouse]
+Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening.
+Thought 2: The paragraph does not tell who Milhouse is named after, maybe I can look up "named after".
+Action 2: Lookup[named after]
+Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous. 
+Thought 3: Milhouse was named after U.S. president Richard Nixon, so the answer is Richard Nixon.
+Action 3: Finish[Richard Nixon]
+Question: Which documentary is about Finnish rock groups, Adam Clayton Powell or The Saimaa Gesture?
+Thought 1: I need to search Adam Clayton Powell and The Saimaa Gesture, and find which documentary is about Finnish rock groups.
+Action 1: Search[Adam Clayton Powell]
+Observation 1: Could not find [Adam Clayton Powell]. Similar: ['Adam Clayton Powell III', 'Seventh Avenue (Manhattan)', 'Adam Clayton Powell Jr. State Office Building', 'Isabel Washington Powell', 'Adam Powell', 'Adam Clayton Powell (film)', 'Giancarlo Esposito'].
+Thought 2: To find the documentary, I can search Adam Clayton Powell (film).
+Action 2: Search[Adam Clayton Powell (film)]
+Observation 2: Adam Clayton Powell is a 1989 American documentary film directed by Richard Kilberg.
+The film is about the rise and fall of influential African-American politician Adam Clayton Powell Jr.[3][4] It was later aired as part of the PBS series The American Experience.
+Thought 3: Adam Clayton Powell (film) is a documentary about an African-American politician, not Finnish rock groups. So the documentary about Finnish rock groups must instead be The Saimaa Gesture.
+Action 3: Finish[The Saimaa Gesture]
+Question: What profession does Nicholas Ray and Elia Kazan have in common?
+Thought 1: I need to search Nicholas Ray and Elia Kazan, find their professions, then find the profession they have in common.
+Action 1: Search[Nicholas Ray]
+Observation 1: Nicholas Ray (born Raymond Nicholas Kienzle Jr., August 7, 1911 – June 16, 1979) was an American film director, screenwriter, and actor best known for the 1955 film Rebel Without a Cause.
+Thought 2: Professions of Nicholas Ray are director, screenwriter, and actor. I need to search Elia Kazan next and find his professions.
+Action 2: Search[Elia Kazan]
+Observation 2: Elia Kazan was an American film and theatre director, producer, screenwriter and actor.
+Thought 3: Professions of Elia Kazan are director, producer, screenwriter, and actor. So profession Nicholas Ray and Elia Kazan have in common is director, screenwriter, and actor.
+Action 3: Finish[director, screenwriter, actor]
+Question: Which magazine was started first Arthur's Magazine or First for Women?
+Thought 1: I need to search Arthur's Magazine and First for Women, and find which was started first.
+Action 1: Search[Arthur's Magazine]
+Observation 1: Arthur's Magazine (1844-1846) was an American literary periodical published in Philadelphia in the 19th century. 
+Thought 2: Arthur's Magazine was started in 1844. I need to search First for Women next.
+Action 2: Search[First for Women]
+Observation 2: First for Women is a woman's magazine published by Bauer Media Group in the USA.[1] The magazine was started in 1989. 
+Thought 3: First for Women was started in 1989. 1844 (Arthur's Magazine) < 1989 (First for Women), so Arthur's Magazine was started first.
+Action 3: Finish[Arthur's Magazine]
+Question: Were Pavel Urysohn and Leonid Levin known for the same type of work?
+Thought 1: I need to search Pavel Urysohn and Leonid Levin, find their types of work, then find if they are the same.
+Action 1: Search[Pavel Urysohn]
+Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory.
+Thought 2: Pavel Urysohn is a mathematician. I need to search Leonid Levin next and find its type of work.
+Action 2: Search[Leonid Levin]
+Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist. 
+Thought 3: Leonid Levin is a mathematician and computer scientist. So Pavel Urysohn and Leonid Levin have the same type of work. 
+Action 3: Finish[yes]
+""" + question)
+    for i in range(1, len(triplets) + 2):
+        s += "Thought " + str(i) + ":"
+        ss = s.fork(1)
+        ss[0] += sgl.gen(name="thought_action", max_tokens=200, stop="Observation")
+        # ss.join()
+        # to verify the correctness of output, this should be collected
+        # print(ss[0]["thought_action"])
+        if i > len(triplets):
+            break
+        s += (triplets[i - 1]["thought"] + "\nAction " + str(i) + ":" +
+              triplets[i - 1]["action"] + "\nObservation " + str(i) + ":" +
+              triplets[i - 1]["observation"] + "\n")
+
+
+def main(args):
+    lines = read_jsonl(args.data_path)[:args.num_questions]
+    arguments = [{
+        "question": k,
+        "triplets": v
+    } for l in lines for k, v in l.items()]
+
+    # Select backend
+    backend = select_sglang_backend(args)
+    sgl.set_default_backend(backend)
+
+    states = []
+    tic = time.time()
+    states = webthink.run_batch(arguments,
+                                temperature=0,
+                                num_threads=args.parallel)
+    latency = time.time() - tic
+
+    # Compute accuracy
+    print(f"Latency: {latency:.3f}")
+
+    # Write results
+    dump_state_text(f"tmp_output_{args.backend}.txt", states)
+
+    with open(args.result_file, "a") as fout:
+        value = {
+            "task": "ReAct Agents",
+            "backend": args.backend,
+            "num_gpus": 1,
+            "latency": round(latency, 3),
+            "num_requests": len(arguments),
+            "other": {
+                "num_questions": args.num_questions,
+                "parallel": args.parallel,
+            },
+        }
+        fout.write(json.dumps(value) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-path", type=str, default="hotpotqa_100.jsonl")
+    parser.add_argument("--num-questions", type=int, default=10)
+    args = add_common_sglang_args_and_parse(parser)
+    main(args)

benchmark/tip_suggestion/README.md

+## Run benchmark
+
+### Benchmark sglang
+```
+python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000
+```
+
+```
+python3 bench_sglang.py --num-questions 64
+python3 bench_sglang.py --num-questions 32 --parallel 1
+```
+
+
+### Benchmark vllm
+```
+python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000
+```
+
+```
+python3 bench_other.py --backend vllm --num-questions 64
+```
+
+
+### Benchmark guidance
+```
+python3 bench_other.py --backend guidance --num-questions 32 --parallel 1
+```

benchmark/tip_suggestion/bench_other.py

+import argparse
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+import json
+import time
+
+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
+
+
+number = 5
+
+
+def expand_tip(topic, tip, generate):
+    s = (
+"""Please expand a tip for a topic into a detailed paragraph.
+
+Topic: staying healthy
+Tip: Regular Exercise
+Paragraph: Incorporate physical activity into your daily routine. This doesn't necessarily mean intense gym workouts; it can be as simple as walking, cycling, or yoga. Regular exercise helps in maintaining a healthy weight, improves cardiovascular health, boosts mental health, and can enhance cognitive function, which is crucial for fields that require intense intellectual engagement.
+
+Topic: building a campfire
+Tip: Choose the Right Location
+Paragraph: Always build your campfire in a safe spot. This means selecting a location that's away from trees, bushes, and other flammable materials. Ideally, use a fire ring if available. If you're building a fire pit, it should be on bare soil or on a bed of stones, not on grass or near roots which can catch fire underground. Make sure the area above is clear of low-hanging branches.
+
+Topic: writing a blog post
+Tip: structure your content effectively
+Paragraph: A well-structured post is easier to read and more enjoyable. Start with an engaging introduction that hooks the reader and clearly states the purpose of your post. Use headings and subheadings to break up the text and guide readers through your content. Bullet points and numbered lists can make information more digestible. Ensure each paragraph flows logically into the next, and conclude with a summary or call-to-action that encourages reader engagement.
+
+Topic: """ + topic + "\nTip: " + tip + "\nParagraph:")
+    return generate(s, max_tokens=128, stop=["\n\n"])
+
+
+def suggest_tips(topic, generate):
+    s = "Please act as a helpful assistant. Your job is to provide users with useful tips on a specific topic.\n"
+    s += "USER: Give some tips for " + topic + ".\n"
+    s += ("ASSISTANT: Okay. Here are " + str(number) + " concise tips, each under 8 words:\n")
+
+    tips = []
+    for i in range(1, 1 + number):
+        s += f"{i}."
+        tip = generate(s, max_tokens=24, stop=[".", "\n"])
+        s += tip + ".\n"
+        tips.append(tip)
+
+    paragraphs = [expand_tip(topic, tip, generate=generate) for tip in tips]
+
+    for i in range(1, 1 + number):
+        s += f"Tip {i}:" + paragraphs[i-1] + "\n"
+    return s
+
+
+def main(args):
+    lines = read_jsonl(args.data_path)[:args.num_questions]
+    states = [None] * len(lines)
+
+    # Select backend
+    if args.backend == "lightllm":
+        url = f"{args.host}:{args.port}/generate"
+        generate = partial(call_generate_lightllm, url=url, temperature=0)
+    elif args.backend == "vllm":
+        url = f"{args.host}:{args.port}/generate"
+        generate = partial(call_generate_vllm, url=url, temperature=0)
+    elif args.backend == "srt-raw":
+        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
+
+        model = models.LlamaCpp("/home/ubuntu/model_weights/Llama-2-7b-chat.gguf", n_gpu_layers=-1, n_ctx=4096)
+
+        def generate(prompt, max_tokens, stop):
+            out = model + prompt + gen(name="answer",
+                max_tokens=max_tokens, temperature=0, stop=stop)
+            return out["answer"]
+
+        # warmup
+        generate("Hello!", max_tokens=8, stop=None)
+    else:
+        raise ValueError(f"Invalid backend: {args.backend}")
+
+    # Run requests
+    def get_one_answer(i):
+        states[i] = suggest_tips(lines[i]["topic"], generate)
+
+    tic = time.time()
+    if args.parallel == 1:
+        for i in tqdm(range(len(lines))):
+            get_one_answer(i)
+    else:
+        with ThreadPoolExecutor(args.parallel) as executor:
+            executor.map(get_one_answer, list(range(len(lines))))
+    latency = time.time() - tic
+
+    # Compute accuracy
+    print(f"Latency: {latency:.3f}")
+
+    # Write results
+    dump_state_text(f"tmp_output_{args.backend}.txt", states)
+
+    with open(args.result_file, "a") as fout:
+        value = {
+            "task": "tip_suggestion",
+            "backend": args.backend,
+            "num_gpus": 1,
+            "latency": round(latency, 3),
+            "num_requests": args.num_questions,
+            "other": {
+                "num_questions": args.num_questions,
+                "parallel": args.parallel,
+            }
+        }
+        fout.write(json.dumps(value) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-path", type=str, default="topic.jsonl")
+    parser.add_argument("--num-questions", type=int, default=100)
+    args = add_common_other_args_and_parse(parser)
+    main(args)

benchmark/tip_suggestion/bench_sglang.py

+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
+
+
+number = 5
+
+
+@sgl.function
+def expand_tip(s, topic, tip):
+    s += (
+"""Please expand a tip for a topic into a detailed paragraph.
+
+Topic: staying healthy
+Tip: Regular Exercise
+Paragraph: Incorporate physical activity into your daily routine. This doesn't necessarily mean intense gym workouts; it can be as simple as walking, cycling, or yoga. Regular exercise helps in maintaining a healthy weight, improves cardiovascular health, boosts mental health, and can enhance cognitive function, which is crucial for fields that require intense intellectual engagement.
+
+Topic: building a campfire
+Tip: Choose the Right Location
+Paragraph: Always build your campfire in a safe spot. This means selecting a location that's away from trees, bushes, and other flammable materials. Ideally, use a fire ring if available. If you're building a fire pit, it should be on bare soil or on a bed of stones, not on grass or near roots which can catch fire underground. Make sure the area above is clear of low-hanging branches.
+
+Topic: writing a blog post
+Tip: structure your content effectively
+Paragraph: A well-structured post is easier to read and more enjoyable. Start with an engaging introduction that hooks the reader and clearly states the purpose of your post. Use headings and subheadings to break up the text and guide readers through your content. Bullet points and numbered lists can make information more digestible. Ensure each paragraph flows logically into the next, and conclude with a summary or call-to-action that encourages reader engagement.
+
+Topic: """ + topic + "\nTip: " + tip + "\nParagraph:")
+    s += sgl.gen("paragraph", max_tokens=128, stop=["\n\n"], temperature=0)
+
+
+@sgl.function
+def suggest_tips(s, topic):
+    s += "Please act as a helpful assistant. Your job is to provide users with useful tips on a specific topic.\n"
+    s += "USER: Give some tips for " + topic + ".\n"
+    s += ("ASSISTANT: Okay. Here are " + str(number) + " concise tips, each under 8 words:\n")
+
+    paragraphs = []
+    for i in range(1, 1 + number):
+        s += f"{i}." + sgl.gen(f"tip_{i}", max_tokens=24, stop=[".", "\n"]) + ".\n"
+        paragraphs.append(expand_tip(topic=topic, tip=s[f"tip_{i}"]))
+
+    for i in range(1, 1 + number):
+        s += f"Tip {i}:" + paragraphs[i-1]["paragraph"] + "\n"
+
+
+def main(args):
+    lines = read_jsonl(args.data_path)[:args.num_questions]
+    arguments = [
+        {"topic": l["topic"]} for l in lines
+    ]
+
+    # Select backend
+    sgl.set_default_backend(select_sglang_backend(args))
+
+    # Run requests
+    tic = time.time()
+    states = suggest_tips.run_batch(
+        arguments, temperature=0, num_threads=args.parallel)
+    latency = time.time() - tic
+
+    # Compute accuracy
+    print(f"Latency: {latency:.3f}")
+
+    # Write results
+    dump_state_text(f"tmp_output_{args.backend}.txt", states)
+
+    with open(args.result_file, "a") as fout:
+        value = {
+            "task": "tip_suggestion",
+            "backend": args.backend,
+            "num_gpus": 1,
+            "latency": round(latency, 3),
+            "num_requests": args.num_questions,
+            "other": {
+                "num_questions": args.num_questions,
+                "parallel": args.parallel,
+            }
+        }
+        fout.write(json.dumps(value) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-path", type=str, default="topic.jsonl")
+    parser.add_argument("--num-questions", type=int, default=100)
+    args = add_common_sglang_args_and_parse(parser)
+    main(args)

benchmark/tree_of_thought/README.md

+## Download data
+```
+wget https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl
+```
+
+## Run benchmark
+
+### Benchmark sglang
+```
+python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000
+```
+
+```
+python3 bench_sglang.py --num-questions 32 --parallel 16
+python3 bench_sglang.py --num-questions 10 --parallel 1
+```
+
+
+### Benchmark vllm
+```
+python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000
+```
+
+```
+python3 bench_other.py --num-questions 32 --backend vllm
+```
+
+
+### Benchmark lightllm
+```
+# A10G
+python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000
+```
+
+```
+python3 bench_other.py --num-questions 32 --backend lightllm
+```
+
+
+### Benchmark guidance
+```
+python3 bench_other.py --num-questions 32 --backend guidance --parallel 1
+```

benchmark/tree_of_thought/bench_other.py

+import argparse
+import ast
+import asyncio
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+import json
+import re
+import time
+
+import numpy as np
+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
+from sglang.utils import read_jsonl, dump_state_text
+
+
+INVALID = -9999999
+
+
+def get_answer_value(answer_str):
+    answer_str = answer_str.replace(",", "")
+    numbers = re.findall(r'\d+', answer_str)
+    if len(numbers) < 1:
+        return INVALID
+    try:
+        return ast.literal_eval(numbers[-1])
+    except SyntaxError:
+        return INVALID
+
+
+def most_frequent_number(numbers):
+    if not numbers:
+        return None
+
+    frequency = Counter(numbers)
+    most_frequent = max(frequency, key=frequency.get)
+    return most_frequent
+
+
+USER_PREFIX = "[INST] "
+USER_SUFFIX = " [/INST]"
+ASSISTANT_PREFIX = ""
+ASSISTANT_SUFFIX = " </s><s>"
+
+# Use a low temp to make the results more deterministic and the comparison more fair.
+temp = 0.3
+
+
+def propose_plan(s, question, num_branches, call_generate):
+    s += (USER_PREFIX +
+"""Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + question + USER_SUFFIX)
+
+    s += ASSISTANT_PREFIX
+    comps = call_generate(s, max_tokens=256, temperature=temp, stop=None, n=num_branches)
+    return [s + comp + ASSISTANT_SUFFIX for comp in comps]
+
+
+def execute_plan(s, num_branches, call_generate):
+    s += (USER_PREFIX +
+"""The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + USER_SUFFIX)
+    s += ASSISTANT_PREFIX
+    comps = call_generate(s, max_tokens=256, temperature=temp, stop=None, n=num_branches)
+    return [s + comp + ASSISTANT_SUFFIX for comp in comps]
+
+
+def reflect_solution(s, num_branches, call_generate):
+    s += (USER_PREFIX +
+"""Okay. Now you evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + USER_SUFFIX)
+    s += ASSISTANT_PREFIX
+    comps = call_generate(s, max_tokens=256, temperature=temp, stop=None, n=num_branches)
+    return [s + comp + ASSISTANT_SUFFIX for comp in comps]
+
+
+def tree_search(question, num_branches, call_generate):
+    s = ""
+    solutions = []
+
+    plan_forks = propose_plan(s, question, num_branches, call_generate)
+    for plan in plan_forks:
+        sol_forks = execute_plan(plan, num_branches, call_generate)
+        for sol in sol_forks:
+            score_forks = reflect_solution(sol, num_branches, call_generate)
+        solutions.append(sol_forks)
+
+    return solutions
+
+
+def main(args):
+    lines = read_jsonl(args.data_path)
+
+    # Construct prompts
+    num_branches = 3
+    questions = []
+    labels = []
+    for i in range(len(lines[:args.num_questions])):
+        questions.append(lines[i]["question"])
+        labels.append(get_answer_value(lines[i]["answer"]))
+    assert all(l != INVALID for l in labels)
+    arguments = [{"question": q, "num_branches": num_branches} for q in questions]
+
+    # Select backend
+    if args.backend == "lightllm":
+        url = f"{args.host}:{args.port}/generate"
+        call_generate = partial(call_generate_lightllm, url=url)
+    elif args.backend == "vllm":
+        url = f"{args.host}:{args.port}/generate"
+        call_generate = partial(call_generate_vllm, url=url)
+    elif args.backend == "srt-raw":
+        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
+
+        model = models.LlamaCpp("/home/ubuntu/model_weights/Llama-2-7b-chat.gguf", n_gpu_layers=-1, n_ctx=4096)
+
+        def call_generate(prompt, temperature, max_tokens, stop, n):
+            if n == 1:
+                out = model + prompt + gen(name="answer",
+                    max_tokens=max_tokens, temperature=temperature, stop=stop)
+                return out["answer"]
+            else:
+                rets = []
+                for i in range(n):
+                    out = model + prompt + gen(name="answer",
+                        max_tokens=max_tokens, temperature=temperature, stop=stop)
+                    rets.append(out["answer"])
+                return rets
+
+    # Run requests
+    states = [None] * len(questions)
+    def get_one_answer(i):
+        states[i] = tree_search(**arguments[i], call_generate=call_generate)
+
+    tic = time.time()
+    if args.parallel == 1:
+        for i in tqdm(range(len(questions))):
+            get_one_answer(i)
+    else:
+        with ThreadPoolExecutor(args.parallel) as executor:
+            executor.map(get_one_answer, list(range(len(questions))))
+    latency = time.time() - tic
+
+    answers_text = []
+    for s in states:
+        answers_text.append([x for xs in s for x in xs])
+
+    preds = []
+    for i in range(len(states)):
+        answers = [get_answer_value(v) for v in answers_text[i]]
+        preds.append(most_frequent_number(answers))
+
+    # Compute accuracy
+    acc = np.mean(np.array(preds) == np.array(labels))
+    invalid = np.mean(np.array(preds) == INVALID)
+    print(f"Latency: {latency:.3f}")
+    print(f"Invalid: {invalid:.3f}")
+    print(f"Accuracy: {acc:.3f}")
+
+    # Write results
+    dump_state_text(f"tmp_output_{args.backend}.txt", answers_text)
+
+    with open(args.result_file, "a") as fout:
+        value = {
+            "task": "tree_of_thought_gsm8k",
+            "backend": args.backend,
+            "num_gpus": 1,
+            "latency": round(latency, 3),
+            "accuracy": round(acc, 3),
+            "num_requests": args.num_questions,
+            "other": {
+                "num_questions": args.num_questions,
+                "parallel": args.parallel,
+            }
+        }
+        fout.write(json.dumps(value) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-path", type=str, default="test.jsonl")
+    parser.add_argument("--num-questions", type=int, default=200)
+    args = add_common_other_args_and_parse(parser)
+    main(args)

benchmark/tree_of_thought/bench_sglang.py

+import argparse
+import ast
+from collections import Counter
+import json
+import re
+import time
+
+import numpy as np
+from sglang.test.test_utils import add_common_sglang_args_and_parse, select_sglang_backend
+from sglang.utils import read_jsonl, dump_state_text
+import sglang as sgl
+
+
+INVALID = -9999999
+
+
+def get_answer_value(answer_str):
+    answer_str = answer_str.replace(",", "")
+    numbers = re.findall(r'\d+', answer_str)
+    if len(numbers) < 1:
+        return INVALID
+    try:
+        return ast.literal_eval(numbers[-1])
+    except SyntaxError:
+        return INVALID
+
+
+def most_frequent_number(numbers):
+    if not numbers:
+        return None
+
+    frequency = Counter(numbers)
+    most_frequent = max(frequency, key=frequency.get)
+    return most_frequent
+
+
+# Use a low temp to make the results more deterministic and the comparison more fair.
+temp = 0.3
+
+
+def propose_plan(s, question, num_branches):
+    s += sgl.user(
+"""Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + question)
+    forks = s.fork(num_branches)
+    forks += sgl.assistant(sgl.gen("plan", max_tokens=256, temperature=temp))
+    return forks
+
+
+def execute_plan(s, num_branches):
+    s += sgl.user(
+"""The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""")
+    forks = s.fork(num_branches)
+    forks += sgl.assistant(sgl.gen("answer", max_tokens=256, temperature=temp))
+    return forks
+
+
+def reflect_solution(s, num_branches):
+    s += sgl.user(
+"""Okay. Now you evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""")
+    forks = s.fork(num_branches)
+    forks += sgl.assistant(sgl.gen("score", max_tokens=256, temperature=temp))
+    return forks
+
+
+@sgl.function
+def tree_search(s, question, num_branches):
+    forks_to_join = []
+
+    plan_forks = propose_plan(s, question, num_branches)
+    forks_to_join.append(plan_forks)
+
+    sol_states = []
+    for plan in plan_forks:
+        forks = execute_plan(plan, num_branches)
+        forks_to_join.append(forks)
+        sol_states.extend(forks)
+
+    for sol in sol_states:
+        forks = reflect_solution(sol, num_branches)
+        forks_to_join.append(forks)
+
+    for f in reversed(forks_to_join):
+        f.join()
+
+
+def main(args):
+    lines = read_jsonl(args.data_path)
+
+    # Construct prompts
+    num_branches = 3
+    questions = []
+    labels = []
+    for i in range(len(lines[:args.num_questions])):
+        questions.append(lines[i]["question"])
+        labels.append(get_answer_value(lines[i]["answer"]))
+    assert all(l != INVALID for l in labels)
+    arguments = [{"question": q, "num_branches": num_branches} for q in questions]
+
+    # Select backend
+    backend = select_sglang_backend(args)
+
+    # Run requests
+    tic = time.time()
+    states = tree_search.run_batch(
+        arguments, temperature=0, backend=backend, num_threads=args.parallel)
+    latency = time.time() - tic
+    answers_text = []
+    for s in states:
+        answers_text.append([x for xs in s["answer"] for x in xs])
+
+    preds = []
+    for i in range(len(states)):
+        answers = [get_answer_value(v) for v in answers_text[i]]
+        preds.append(most_frequent_number(answers))
+
+    # Compute accuracy
+    acc = np.mean(np.array(preds) == np.array(labels))
+    invalid = np.mean(np.array(preds) == INVALID)
+    print(f"Latency: {latency:.3f}")
+    print(f"Invalid: {invalid:.3f}")
+    print(f"Accuracy: {acc:.3f}")
+
+    # Write results
+    dump_state_text(f"tmp_output_{args.backend}.txt", answers_text)
+
+    with open(args.result_file, "a") as fout:
+        value = {
+            "task": "tree_of_thought_gsm8k",
+            "backend": args.backend,
+            "num_gpus": 1,
+            "latency": round(latency, 3),
+            "accuracy": round(acc, 3),
+            "num_requests": args.num_questions,
+            "other": {
+                "num_questions": args.num_questions,
+                "parallel": args.parallel,
+            }
+        }
+        fout.write(json.dumps(value) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-path", type=str, default="test.jsonl")
+    parser.add_argument("--num-questions", type=int, default=200)
+    args = add_common_sglang_args_and_parse(parser)
+    main(args)

benchmark/tree_of_thought_deep/README.md

+## Download data
+```
+wget https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl
+```
+
+## Run benchmark
+
+### Benchmark sglang
+```
+python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000
+```
+
+```
+python3 bench_sglang.py --num-questions 32 --parallel 8
+python3 bench_sglang.py --num-questions 16 --parallel 1
+```
+
+
+### Benchmark vllm
+```
+python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000
+```
+
+```
+python3 bench_other.py --num-questions 32 --backend vllm
+```
+
+
+### Benchmark lightllm
+```
+# A10G
+python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000
+```
+
+```
+python3 bench_other.py --num-questions 32 --backend lightllm
+```
+
+
+### Benchmark guidance
+```
+python3 bench_other.py --num-questions 8 --backend guidance --parallel 1
+```

benchmark/tree_of_thought_deep/bench_other.py

+import argparse
+import ast
+import asyncio
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+import json
+import re
+import time
+
+import numpy as np
+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
+from sglang.utils import read_jsonl, dump_state_text
+
+
+INVALID = -9999999
+
+
+def get_answer_value(answer_str):
+    answer_str = answer_str.replace(",", "")
+    numbers = re.findall(r'\d+', answer_str)
+    if len(numbers) < 1:
+        return INVALID
+    try:
+        return ast.literal_eval(numbers[-1])
+    except SyntaxError:
+        return INVALID
+
+
+def most_frequent_number(numbers):
+    if not numbers:
+        return None
+
+    frequency = Counter(numbers)
+    most_frequent = max(frequency, key=frequency.get)
+    return most_frequent
+
+
+USER_PREFIX = "[INST] "
+USER_SUFFIX = " [/INST]"
+ASSISTANT_PREFIX = ""
+ASSISTANT_SUFFIX = " </s><s>"
+
+# Use a low temp to make the results more deterministic and the comparison more fair.
+temp = 0.001
+
+
+def propose_plan(s, question, num_branches, call_generate):
+    s += (USER_PREFIX +
+"""Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + question + USER_SUFFIX)
+
+    s += ASSISTANT_PREFIX
+    comps = call_generate(s, max_tokens=256, temperature=temp, stop=None, n=num_branches)
+    return [s + comp + ASSISTANT_SUFFIX for comp in comps]
+
+
+def execute_plan(s, num_branches, call_generate):
+    s += (USER_PREFIX +
+"""The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + USER_SUFFIX)
+    s += ASSISTANT_PREFIX
+    comps = call_generate(s, max_tokens=256, temperature=temp, stop=None, n=num_branches)
+    return [s + comp + ASSISTANT_SUFFIX for comp in comps]
+
+
+def reflect_solution(s, num_branches, call_generate):
+    s += (USER_PREFIX +
+"""Okay. Now, evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + USER_SUFFIX)
+    s += ASSISTANT_PREFIX
+    comps = call_generate(s, max_tokens=256, temperature=temp, stop=None, n=num_branches)
+    return [s + comp + ASSISTANT_SUFFIX for comp in comps]
+
+
+def get_final_answer(s, num_branches, call_generate):
+    s += (USER_PREFIX +
+"""Based on your reflection, do you change your mind? Now, give me the final answer after careful consideration.""" + USER_SUFFIX)
+    s += ASSISTANT_PREFIX
+    comps = call_generate(s, max_tokens=256, temperature=temp, stop=None, n=num_branches)
+    return [s + comp + ASSISTANT_SUFFIX for comp in comps]
+
+
+def tree_search(question, num_branches, call_generate):
+    plan_forks = propose_plan("", question, num_branches, call_generate)
+
+    sol_states = []
+    for plan in plan_forks:
+        forks = execute_plan(plan, num_branches, call_generate)
+        sol_states.extend(forks)
+
+    ref_states = []
+    for sol in sol_states:
+        forks = reflect_solution(sol, num_branches, call_generate)
+        ref_states.extend(forks)
+
+    solutions = []
+    for sol in ref_states:
+        ans = get_final_answer(sol, num_branches, call_generate)
+        solutions.append(ans)
+
+    return solutions
+
+
+def main(args):
+    lines = read_jsonl(args.data_path)
+
+    # Construct prompts
+    num_branches = 2
+    questions = []
+    labels = []
+    for i in range(len(lines[:args.num_questions])):
+        questions.append(lines[i]["question"])
+        labels.append(get_answer_value(lines[i]["answer"]))
+    assert all(l != INVALID for l in labels)
+    arguments = [{"question": q, "num_branches": num_branches} for q in questions]
+
+    # Select backend
+    if args.backend == "lightllm":
+        url = f"{args.host}:{args.port}/generate"
+        call_generate = partial(call_generate_lightllm, url=url)
+    elif args.backend == "vllm":
+        url = f"{args.host}:{args.port}/generate"
+        call_generate = partial(call_generate_vllm, url=url)
+    elif args.backend == "srt-raw":
+        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
+
+        model = models.LlamaCpp("/home/ubuntu/model_weights/Llama-2-7b-chat.gguf", n_gpu_layers=-1, n_ctx=4096)
+
+        def call_generate(prompt, temperature, max_tokens, stop, n):
+            if n == 1:
+                out = model + prompt + gen(name="answer",
+                    max_tokens=max_tokens, temperature=temperature, stop=stop)
+                return out["answer"]
+            else:
+                rets = []
+                for i in range(n):
+                    out = model + prompt + gen(name="answer",
+                        max_tokens=max_tokens, temperature=temperature, stop=stop)
+                    rets.append(out["answer"])
+                return rets
+
+    # Run requests
+    states = [None] * len(questions)
+    def get_one_answer(i):
+        states[i] = tree_search(**arguments[i], call_generate=call_generate)
+
+    tic = time.time()
+    if args.parallel == 1:
+        for i in tqdm(range(len(questions))):
+            get_one_answer(i)
+    else:
+        with ThreadPoolExecutor(args.parallel) as executor:
+            executor.map(get_one_answer, list(range(len(questions))))
+    latency = time.time() - tic
+
+    answers_text = []
+    for s in states:
+        answers_text.append([x for xs in s for x in xs])
+
+    preds = []
+    for i in range(len(states)):
+        answers = [get_answer_value(v) for v in answers_text[i]]
+        preds.append(most_frequent_number(answers))
+
+    # Compute accuracy
+    acc = np.mean(np.array(preds) == np.array(labels))
+    invalid = np.mean(np.array(preds) == INVALID)
+    print(f"Latency: {latency:.3f}")
+    print(f"Invalid: {invalid:.3f}")
+    print(f"Accuracy: {acc:.3f}")
+
+    # Write results
+    dump_state_text(f"tmp_output_{args.backend}.txt", answers_text)
+
+    with open(args.result_file, "a") as fout:
+        value = {
+            "task": "tree_of_thought_gsm8k",
+            "backend": args.backend,
+            "num_gpus": 1,
+            "latency": round(latency, 3),
+            "accuracy": round(acc, 3),
+            "num_requests": args.num_questions,
+            "other": {
+                "num_questions": args.num_questions,
+                "parallel": args.parallel,
+            }
+        }
+        fout.write(json.dumps(value) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-path", type=str, default="test.jsonl")
+    parser.add_argument("--num-questions", type=int, default=200)
+    args = add_common_other_args_and_parse(parser)
+    main(args)

benchmark/tree_of_thought_deep/bench_sglang.py

+import argparse
+import ast
+from collections import Counter
+import json
+import re
+import time
+
+import numpy as np
+from sglang.test.test_utils import add_common_sglang_args_and_parse, select_sglang_backend
+from sglang.utils import read_jsonl, dump_state_text
+import sglang as sgl
+
+
+INVALID = -9999999
+
+
+def get_answer_value(answer_str):
+    answer_str = answer_str.replace(",", "")
+    numbers = re.findall(r'\d+', answer_str)
+    if len(numbers) < 1:
+        return INVALID
+    try:
+        return ast.literal_eval(numbers[-1])
+    except SyntaxError:
+        return INVALID
+
+
+def most_frequent_number(numbers):
+    if not numbers:
+        return None
+
+    frequency = Counter(numbers)
+    most_frequent = max(frequency, key=frequency.get)
+    return most_frequent
+
+
+# Use a low temp to make the results more deterministic and the comparison more fair.
+temp = 0.001
+
+
+def propose_plan(s, question, num_branches):
+    s += sgl.user(
+"""Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + question)
+    forks = s.fork(num_branches)
+    forks += sgl.assistant(sgl.gen("plan", max_tokens=256, temperature=temp))
+    return forks
+
+
+def execute_plan(s, num_branches):
+    s += sgl.user(
+"""The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""")
+    forks = s.fork(num_branches)
+    forks += sgl.assistant(sgl.gen("answer", max_tokens=256, temperature=temp))
+    return forks
+
+
+def reflect_solution(s, num_branches):
+    s += sgl.user(
+"""Okay. Now, evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""")
+    forks = s.fork(num_branches)
+    forks += sgl.assistant(sgl.gen("score", max_tokens=256, temperature=temp))
+    return forks
+
+
+def get_final_answer(s, num_branches):
+    s += sgl.user(
+"""Based on your reflection, do you change your mind? Now, give me the final answer after careful consideration.""")
+    forks = s.fork(num_branches)
+    forks += sgl.assistant(sgl.gen("final_answer", max_tokens=256, temperature=temp))
+    return forks
+
+
+
+@sgl.function
+def tree_search(s, question, num_branches):
+    plan_forks = propose_plan(s, question, num_branches)
+
+    sol_states = []
+    for plan in plan_forks:
+        forks = execute_plan(plan, num_branches)
+        sol_states.extend(forks)
+
+    ref_states = []
+    for sol in sol_states:
+        forks = reflect_solution(sol, num_branches)
+        ref_states.extend(forks)
+
+    solutions = []
+    for sol in ref_states:
+        forks = get_final_answer(sol, num_branches)
+        solutions.append(forks)
+    solutions = [[s.text() for s in forks] for forks in solutions]
+
+    return solutions
+
+def main(args):
+    lines = read_jsonl(args.data_path)
+
+    # Construct prompts
+    num_branches = 2
+    questions = []
+    labels = []
+    for i in range(len(lines[:args.num_questions])):
+        questions.append(lines[i]["question"])
+        labels.append(get_answer_value(lines[i]["answer"]))
+    assert all(l != INVALID for l in labels)
+    arguments = [{"question": q, "num_branches": num_branches} for q in questions]
+
+    # Select backend
+    backend = select_sglang_backend(args)
+
+    # Run requests
+    tic = time.time()
+    states = tree_search.run_batch(
+        arguments, temperature=0, backend=backend, num_threads=args.parallel)
+    latency = time.time() - tic
+    answers_text = []
+    for s in states:
+        answers_text.append([x for xs in s.ret_value for x in xs])
+
+    preds = []
+    for i in range(len(states)):
+        answers = [get_answer_value(v) for v in answers_text[i]]
+        preds.append(most_frequent_number(answers))
+
+    # Compute accuracy
+    acc = np.mean(np.array(preds) == np.array(labels))
+    invalid = np.mean(np.array(preds) == INVALID)
+    print(f"Latency: {latency:.3f}")
+    print(f"Invalid: {invalid:.3f}")
+    print(f"Accuracy: {acc:.3f}")
+
+    # Write results
+    dump_state_text(f"tmp_output_{args.backend}.txt", answers_text)
+
+    with open(args.result_file, "a") as fout:
+        value = {
+            "task": "tree_of_thought_gsm8k",
+            "backend": args.backend,
+            "num_gpus": 1,
+            "latency": round(latency, 3),
+            "accuracy": round(acc, 3),
+            "num_requests": args.num_questions,
+            "other": {
+                "num_questions": args.num_questions,
+                "parallel": args.parallel,
+            }
+        }
+        fout.write(json.dumps(value) + "\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-path", type=str, default="test.jsonl")
+    parser.add_argument("--num-questions", type=int, default=200)
+    args = add_common_sglang_args_and_parse(parser)
+    main(args)

docs/flashinfer.md

+## Flashinfer Mode
+
+[`flashinfer`](https://github.com/flashinfer-ai/flashinfer) is a kernel library for LLM serving; we use it here to support our attention computation.
+
+### Install flashinfer
+
+```bash
+git submodule update --init --recursive
+pip install 3rdparty/flashinfer/python
+```
+
+### Run Sever With Flashinfer Mode
+
+Add through `--model_mode` argument from the command line.
+
+Example:
+
+```bash
+python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --model-mode flashinfer
+```
\ No newline at end of file