40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
|
|
|
|
import argparse
|
|
from typing import Dict, List
|
|
import os
|
|
import sglang as sgl
|
|
from sglang import OpenAI, assistant, gen, set_default_backend, system, user
|
|
|
|
|
|
|
|
@sgl.function
|
|
def multi_turn_question(s, question_1, question_2):
|
|
s += sgl.system("You are a helpful assistant.")
|
|
s += sgl.user(question_1)
|
|
s += sgl.assistant(sgl.gen("answer_1", max_tokens=256))
|
|
s += sgl.user(question_2)
|
|
s += sgl.assistant(sgl.gen("answer_2", max_tokens=256))
|
|
|
|
|
|
|
|
def test_sglang():
|
|
print('\n----- Test sglang ---')
|
|
state = multi_turn_question.run(
|
|
question_1="What is the capital of the United States?",
|
|
question_2="List two local attractions.",
|
|
)
|
|
|
|
for m in state.messages():
|
|
print(m["role"], ":", m["content"])
|
|
|
|
print("\n-- answer_1 --\n", state["answer_1"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import time
|
|
start_time = time.time()
|
|
sgl.set_default_backend(sgl.OpenAI("gpt-4"))
|
|
|
|
test_sglang()
|
|
print(f"Time taken: {time.time() - start_time:.2f} seconds, or {(time.time() - start_time) / 60:.2f} minutes, or {(time.time() - start_time) / 3600:.2f} hours.\a") |