From ba893ed1486b41afe32493fbf180cc50dedc282b Mon Sep 17 00:00:00 2001 From: alpayariyak Date: Wed, 20 Dec 2023 18:40:43 -0500 Subject: [PATCH] Refactor prompt handling in handler.py --- src/handler.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/handler.py b/src/handler.py index 7c59d7a..6d5f059 100644 --- a/src/handler.py +++ b/src/handler.py @@ -18,8 +18,10 @@ async def handler(job: dict) -> Generator[dict, None, None]: if messages: prompt = tokenizer.apply_chat_template(messages) - elif apply_chat_template: + elif prompt and apply_chat_template: prompt = tokenizer.apply_chat_template(prompt) + elif not prompt: + raise ValueError("Must specify prompt or messages") streaming = job_input.get("streaming", False) batch_size = job_input.get("batch_size", serverless_config.default_batch_size)