more changes

app.py

@@ -36,30 +36,24 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
         return_tensors="pt",
     ).to("cuda" if torch.cuda.is_available() else "cpu")
 
-    # Use TextStreamer to process and yield outputs incrementally
-    class GradioStreamer(TextStreamer):
-        def __init__(self, tokenizer, *args, **kwargs):
-            super().__init__(tokenizer, *args, **kwargs)
-            self.generated_text = ""
+    attention_mask = inputs.ne(tokenizer.pad_token_id).long()  # Explicitly set attention mask
 
-        def on_token(self, token_id):
-            token = self.tokenizer.decode(token_id, skip_special_tokens=True)
-            self.generated_text += token
-            yield self.generated_text
-
-    # Initialize Gradio-compatible streamer
-    streamer = GradioStreamer(tokenizer, skip_prompt=True)
-    
-    # Generate response with streaming
-    _ = model.generate(
+    # Generate response tokens
+    generated_tokens = model.generate(
         input_ids=inputs,
+        attention_mask=attention_mask,
         max_new_tokens=max_tokens,
         use_cache=True,
         temperature=temperature,
         top_p=top_p,
-        streamer=streamer,
     )
 
+    # Decode generated tokens
+    response = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
+
+    # Yield response in the required Gradio format
+    yield [{"role": "assistant", "content": response}]
+
 
 
 # Define the Gradio interface