owen/worker-vllm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Chat Template

Browse Source
This commit is contained in:
alpayariyak
2023-12-20 18:37:54 -05:00
parent dad02e13d8
commit 582e21f97f
2 changed files with 26 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/handler.py
+10 -2
View File
@@ -5,14 +5,22 @@ from utils import validate_and_convert_sampling_params, initialize_llm_engine, S
from vllm.utils import random_uuid
serverless_config = ServerlessConfig()
llm = initialize_llm_engine()
llm, tokenizer = initialize_llm_engine()
def concurrency_modifier(current_concurrency) -> int:
return max(0, serverless_config.max_concurrency - current_concurrency)
async def handler(job: dict) -> Generator[dict, None, None]:
job_input = job["input"]
prompt = job_input["prompt"]
prompt = job_input.get("prompt")
apply_chat_template = job_input.get("apply_chat_template", False)
messages = job_input.get("messages")
if messages:
prompt = tokenizer.apply_chat_template(messages)
elif apply_chat_template:
prompt = tokenizer.apply_chat_template(prompt)
streaming = job_input.get("streaming", False)
batch_size = job_input.get("batch_size", serverless_config.default_batch_size)
sampling_params = job_input.get("sampling_params", {})