Use background threading for generation

app.py

@@ -5,41 +5,41 @@ import gradio as gr
 import spaces
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 
-# @spaces.GPU
-# def chat_with_model(messages):
-#     global current_model, current_tokenizer
-#     if current_model is None or current_tokenizer is None:
-#         yield messages + [{"role": "assistant", "content": "⚠️ No model loaded."}]
-#         return
+@spaces.GPU
+def chat_with_model(messages):
+    global current_model, current_tokenizer
+    if current_model is None or current_tokenizer is None:
+        yield messages + [{"role": "assistant", "content": "⚠️ No model loaded."}]
+        return
 
-#     current_model.to("cuda").half()
+    current_model.to("cuda").half()
 
-#     prompt = format_prompt(messages)
-#     inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
+    prompt = format_prompt(messages)
+    inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
 
-#     streamer = TextIteratorStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=True)
-#     generation_kwargs = dict(
-#         **inputs,
-#         max_new_tokens=256,
-#         do_sample=True,
-#         streamer=streamer
-#     )
+    streamer = TextIteratorStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = dict(
+        **inputs,
+        max_new_tokens=256,
+        do_sample=True,
+        streamer=streamer
+    )
 
-#     # Launch generation in a background thread
-#     thread = threading.Thread(target=current_model.generate, kwargs=generation_kwargs)
-#     thread.start()
+    # Launch generation in a background thread
+    thread = threading.Thread(target=current_model.generate, kwargs=generation_kwargs)
+    thread.start()
 
-#     output_text = ""
-#     messages = messages.copy()
-#     messages.append({"role": "assistant", "content": ""})
+    output_text = ""
+    messages = messages.copy()
+    messages.append({"role": "assistant", "content": ""})
 
-#     for new_text in streamer:
-#         output_text += new_text
-#         messages[-1]["content"] = output_text
-#         yield messages
+    for new_text in streamer:
+        output_text += new_text
+        messages[-1]["content"] = output_text
+        yield messages
 
-#     current_model.to("cpu")
-#     torch.cuda.empty_cache()
+    current_model.to("cpu")
+    torch.cuda.empty_cache()
 
 
 # Globals
@@ -92,41 +92,41 @@ with gr.Blocks() as demo:
 
     default_model = gr.State("meta-llama/Llama-3.2-3B-Instruct")
 
-    @spaces.GPU
-    def chat_with_model(messages):
-        global current_model, current_tokenizer
-        if current_model is None or current_tokenizer is None:
-            yield messages + [{"role": "assistant", "content": "⚠️ No model loaded."}]
-            return
-
-        current_model = current_model.to("cuda").half()
-
-        prompt = format_prompt(messages)
-        inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
-
-        output_ids = []
-        messages = messages.copy()
-        messages.append({"role": "assistant", "content": ""})
-
-        for token_id in current_model.generate(
-            **inputs,
-            max_new_tokens=256,
-            do_sample=True,
-            return_dict_in_generate=True,
-            output_scores=False
-        ).sequences[0][inputs['input_ids'].shape[-1]:]:  # skip input tokens
-            output_ids.append(token_id.item())
-            decoded = current_tokenizer.decode(output_ids, skip_special_tokens=False)
-            if output_ids[-1] == current_tokenizer.eos_token_id:
-                current_model.to("cpu")
-                torch.cuda.empty_cache()
-                return
-            messages[-1]["content"] = decoded
-            yield messages
-
-        current_model.to("cpu")
-        torch.cuda.empty_cache()
-        return
+    # @spaces.GPU
+    # def chat_with_model(messages):
+    #     global current_model, current_tokenizer
+    #     if current_model is None or current_tokenizer is None:
+    #         yield messages + [{"role": "assistant", "content": "⚠️ No model loaded."}]
+    #         return
+
+    #     current_model = current_model.to("cuda").half()
+
+    #     prompt = format_prompt(messages)
+    #     inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
+
+    #     output_ids = []
+    #     messages = messages.copy()
+    #     messages.append({"role": "assistant", "content": ""})
+
+    #     for token_id in current_model.generate(
+    #         **inputs,
+    #         max_new_tokens=256,
+    #         do_sample=True,
+    #         return_dict_in_generate=True,
+    #         output_scores=False
+    #     ).sequences[0][inputs['input_ids'].shape[-1]:]:  # skip input tokens
+    #         output_ids.append(token_id.item())
+    #         decoded = current_tokenizer.decode(output_ids, skip_special_tokens=False)
+    #         if output_ids[-1] == current_tokenizer.eos_token_id:
+    #             current_model.to("cpu")
+    #             torch.cuda.empty_cache()
+    #             return
+    #         messages[-1]["content"] = decoded
+    #         yield messages
+
+    #     current_model.to("cpu")
+    #     torch.cuda.empty_cache()
+    #     return
 
     with gr.Row():
         model_selector = gr.Dropdown(choices=model_choices, label="Select Model")