Files
sglang/benchmark/llava_bench/bench_sglang.py
Lianmin Zheng 22085081bb 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>
2024-01-08 04:37:50 +00:00

97 lines
3.0 KiB
Python

import argparse
import json
import time
import os
import sglang as sgl
import tqdm
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 PIL import Image
@sgl.function
def image_qa(s, image_file, question):
s += sgl.user(sgl.image(image_file) + question)
s += sgl.assistant(sgl.gen("answer", max_tokens=args.max_tokens))
def main(args):
lines = read_jsonl(args.question_file)[:args.num_questions]
arguments = [
{"image_file":
os.path.abspath(args.image_folder + "/" + l["image"]),
"question": l["text"]} for l in lines
]
#arguments = [
# {"image_file":
# Image.open(os.path.abspath(args.image_folder + "/" + l["image"])),
# "question": l["text"]} for l in lines
#]
states = [None] * len(lines)
# Select backend
backend = select_sglang_backend(args)
sgl.set_default_backend(backend)
# Run requests
tic = time.time()
if args.parallel == 1:
for i in tqdm.tqdm(range(len(lines))):
image_file = arguments[i]["image_file"]
question = arguments[i]["question"]
ret = image_qa.run(
image_file=image_file,
question=question,
temperature=0)
states[i] = ret
else:
states = image_qa.run_batch(
arguments,
temperature=0,
num_threads=args.parallel,
progress_bar=True)
latency = time.time() - tic
print(f"Latency: {latency:.3f}")
# Write results
dump_state_text(f"tmp_output_{args.backend}.txt", states)
print(f"Write output to {args.answer_file}")
with open(args.answer_file, "w") as fout:
for i in range(len(lines)):
value = {
"question_id": lines[i]["question_id"],
"prompt": lines[i]["text"],
"text": states[i]["answer"].strip(),
"model_id": backend.model_info["model_path"],
"answer_id": i,
"metadata": {},
}
fout.write(json.dumps(value) + "\n")
with open(args.result_file, "a") as fout:
value = {
"task": "llava_bench",
"backend": args.backend,
"num_gpus": 1,
"latency": round(latency, 3),
"num_requests": len(lines),
"parallel": args.parallel,
}
fout.write(json.dumps(value) + "\n")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--question-file", type=str, default="questions.jsonl")
parser.add_argument("--answer-file", type=str, default="answers.jsonl")
parser.add_argument("--image-folder", type=str, default="./images")
parser.add_argument("--temperature", type=float, default=0.0)
parser.add_argument("--num-questions", type=int, default=None)
parser.add_argument("--max-tokens", type=int, default=768)
args = add_common_sglang_args_and_parse(parser)
main(args)