Moving spaces.GPU

main.py

@@ -57,6 +57,29 @@ def reformat_math(text):
     return text
 
 
+@spaces.GPU
+def _generate(history):
+    text = tokenizer.apply_chat_template(
+        history,
+        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
+
+
 async def chat(prompt, history):
     """Respond to a chat prompt."""
     message = {
@@ -64,31 +87,12 @@ async def chat(prompt, history):
         "content": prompt,
     }
 
+    # build the messages list
     history = [] if history is None else history
+    message_list = history + [message]
 
-    @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()
+    # get the task and the streamer
+    task, streamer = _generate(message_list)
 
     buffer = ""
     reasoning = ""