Spaces:
Sleeping
Sleeping
Commit ·
792120a
1
Parent(s): a5eb894
Remove repetition stuff
Browse files
main.py
CHANGED
|
@@ -161,7 +161,7 @@ def predict(inp):
|
|
| 161 |
device = next(model.parameters()).device
|
| 162 |
inputs_device = {k: v.to(device) for k, v in inputs.items()}
|
| 163 |
# Call model.generate which will use streamer to yield tokens; model weights stay on their device
|
| 164 |
-
model.generate(**inputs_device, max_new_tokens=new_tokens, streamer=streamer
|
| 165 |
except Exception as e:
|
| 166 |
try:
|
| 167 |
streamer.end()
|
|
|
|
| 161 |
device = next(model.parameters()).device
|
| 162 |
inputs_device = {k: v.to(device) for k, v in inputs.items()}
|
| 163 |
# Call model.generate which will use streamer to yield tokens; model weights stay on their device
|
| 164 |
+
model.generate(**inputs_device, max_new_tokens=new_tokens, streamer=streamer)
|
| 165 |
except Exception as e:
|
| 166 |
try:
|
| 167 |
streamer.end()
|