Spaces:
Sleeping
Sleeping
dntrplytch
commited on
if deivce type is CPU, remove the timeout argument to TextIteratorStreamer
Browse filesthe 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.
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 |
-
|
|
|
|
|
|
|
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,
|