dntrplytch commited on
Commit
28318a9
·
verified ·
1 Parent(s): 6fb91aa

if deivce type is CPU, remove the timeout argument to TextIteratorStreamer

Browse files

the timeout argument to TextIteratorStreamer (was 20 s) failed on the CPU. It seemed to be taking ~25 seconds. Removing it altogether apparently make the wait indefinite.
My test with 120s. as the value also worked without the code erroring out.

Files changed (1) hide show
  1. app.py +4 -1
app.py CHANGED
@@ -57,7 +57,10 @@ def generate(
57
  gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
58
  input_ids = input_ids.to(model.device)
59
 
60
- streamer = TextIteratorStreamer(tokenizer, timeout=120.0, skip_prompt=True, skip_special_tokens=True)
 
 
 
61
  generate_kwargs = dict(
62
  {"input_ids": input_ids},
63
  streamer=streamer,
 
57
  gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
58
  input_ids = input_ids.to(model.device)
59
 
60
+ if device == "cpu":
61
+ streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
62
+ else:
63
+ streamer = TextIteratorStreamer(tokenizer, timeout=120.0, skip_prompt=True, skip_special_tokens=True)
64
  generate_kwargs = dict(
65
  {"input_ids": input_ids},
66
  streamer=streamer,