Update app.py

app.py

@@ -36,12 +36,18 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    # Tokenize the input messages
+    # Tokenize the input messages with dynamic padding and truncation
     input_text = system_message + " ".join([f"{msg['role']}: {msg['content']}" for msg in messages])
-    input_ids = tokenizer.encode(input_text, return_tensors="pt")
+    inputs = tokenizer(
+        input_text, 
+        return_tensors="pt", 
+        padding=True,  # Dynamically pad to the longest sequence in the batch
+        truncation=True,  # Truncate if exceeds max length
+        max_length=max_tokens  # Ensure max length is respected
+    )
 
-    # Create attention mask
-    attention_mask = input_ids.ne(tokenizer.pad_token_id).long()
+    input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
 
     # Generate a response
     chat_history_ids = model.generate(
@@ -51,7 +57,7 @@ def respond(
         top_p=top_p,
         pad_token_id=tokenizer.eos_token_id,
         do_sample=True,
-        attention_mask=attention_mask,
+        attention_mask=attention_mask,  # Use the dynamically generated attention mask
     )
 
     # Decode the response