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