dntrplytch commited on
Commit
1d5e7ac
·
verified ·
1 Parent(s): 9147ca6

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

Browse files

the timeout argument to TextIteratorStreamer (was 20 s) failed to have the program run 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 on the CPU without the code erroring out.

Files changed (1) hide show
  1. app.py +5 -1
app.py CHANGED
@@ -57,7 +57,11 @@ 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=20.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
+
65
  generate_kwargs = dict(
66
  {"input_ids": input_ids},
67
  streamer=streamer,