Skip to content

Commit

Permalink
Expose timeout for stremer for fastchat worker
Browse files Browse the repository at this point in the history
  • Loading branch information
Oscilloscope98 committed Jun 12, 2024
1 parent f224e98 commit fb34023
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion python/llm/src/ipex_llm/serving/fastchat/ipex_llm_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ def __init__(
load_low_bit_model: bool = False,
stream_interval: int = 4,
benchmark: str = "true",
streamer_timeout: int = 60,
):
super().__init__(
controller_addr,
Expand All @@ -83,6 +84,8 @@ def __init__(
conv_template,
)

self.streamer_timeout = streamer_timeout

self.load_in_low_bit = load_in_low_bit
self.load_low_bit_model = load_low_bit_model
logger.info(
Expand Down Expand Up @@ -323,7 +326,7 @@ def generate_stream_gate(self, params):
# Use TextIteratorStreamer for streaming output
streamer = TextIteratorStreamer(
tokenizer=self.tokenizer,
timeout=60,
timeout=self.streamer_timeout,
skip_prompt=True,
skip_special_tokens=True,
)
Expand Down

0 comments on commit fb34023

Please sign in to comment.