Update app.py

app.py

@@ -9,6 +9,9 @@ from transformers import (
     TextIteratorStreamer,
 )
 
+import spaces
+
+
 MODEL_ID = os.getenv("MODEL_ID", "yasserrmd/SoftwareArchitecture-Instruct-v1")
 
 # -------- Load model & tokenizer --------
@@ -56,45 +59,45 @@ def format_history_as_messages(history):
             messages.append({"role": "assistant", "content": a})
     return messages
 
+@spaces.GPU
 def stream_generate(messages, max_new_tokens, temperature, top_p, repetition_penalty, seed=None):
-    """
-    Stream text from model.generate using TextIteratorStreamer.
-    """
     if seed is not None and seed >= 0:
         torch.manual_seed(seed)
 
     inputs = tokenizer.apply_chat_template(
         messages,
-        add_generation_prompt=True,   # IMPORTANT for chat models
+        add_generation_prompt=True,
         return_tensors="pt",
         tokenize=True,
         return_dict=True,
     )
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+
+    # Keep only what the model expects
+    allowed = {"input_ids", "attention_mask"}  # no token_type_ids for causal LMs
+    inputs = {k: v.to(model.device) for k, v in inputs.items() if k in allowed}
 
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     gen_kwargs = dict(
         **inputs,
-        max_new_tokens=max_new_tokens,
+        max_new_tokens=int(max_new_tokens),
         temperature=float(temperature),
         top_p=float(top_p),
         repetition_penalty=float(repetition_penalty),
-        do_sample=True if temperature > 0 else False,
+        do_sample=temperature > 0,
         use_cache=True,
         streamer=streamer,
     )
 
-    # Run generation in a thread so we can yield from streamer
     thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
     thread.start()
 
-    partial_text = ""
-    for new_text in streamer:
-        partial_text += new_text
-        yield partial_text
+    partial = ""
+    for chunk in streamer:
+        partial += chunk
+        yield partial
 
 # -------- Gradio callbacks --------
-
+@spaces.GPU
 def chat_respond(user_msg, chat_history, max_new_tokens, temperature, top_p, repetition_penalty, seed):
     if not user_msg or not user_msg.strip():
         return gr.update(), chat_history