Moving spaces.GPU

main.py

@@ -57,12 +57,6 @@ def reformat_math(text):
     return text
 
 
-@spaces.GPU
-def generate(model, **kwargs):
-    """Geneerate text using the model."""
-    model.generate(**kwargs)
-
-
 async def chat(prompt, history):
     """Respond to a chat prompt."""
     message = {
@@ -71,25 +65,30 @@ async def chat(prompt, history):
     }
 
     history = [] if history is None else history
-    text = tokenizer.apply_chat_template(
-        history + [message],
-        tokenize=False,
-        add_generation_prompt=True,
-    )
 
-    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
-    streamer = AsyncTextIteratorStreamer(tokenizer, skip_special_tokens=True)
-
-    task = asyncio.get_running_loop().run_in_executor(
-        None,
-        functools.partial(
-            generate,
-            model,
-            **model_inputs,
-            max_new_tokens=1024 * 128,
-            streamer=streamer,
-        ),
-    )
+    @spaces.GPU
+    def _generate():
+        text = tokenizer.apply_chat_template(
+            history + [message],
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+
+        model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+        streamer = AsyncTextIteratorStreamer(tokenizer, skip_special_tokens=True)
+
+        task = asyncio.get_running_loop().run_in_executor(
+            None,
+            functools.partial(
+                model.generate,
+                max_new_tokens=1024 * 128,
+                streamer=streamer,
+                **model_inputs,
+            ),
+        )
+        return task, streamer
+
+    task, streamer = _generate()
 
     buffer = ""
     reasoning = ""