Update app.py

app.py

@@ -36,19 +36,16 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    # Tokenize the input messages with dynamic padding and truncation
+    # Tokenize the input messages
     input_text = system_message + " ".join([f"{msg['role']}: {msg['content']}" for msg in messages])
-    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
-    )
-
-    input_ids = inputs["input_ids"]
-    attention_mask = inputs["attention_mask"]
-
+    input_ids = tokenizer.encode(input_text, return_tensors="pt")
+    
+    # Move input_ids to the GPU
+    input_ids = input_ids.to("cuda")
+    
+    # Create attention mask and move to GPU
+    attention_mask = input_ids.ne(tokenizer.pad_token_id).long().to("cuda")
+    
     # Generate a response
     chat_history_ids = model.generate(
         input_ids,
@@ -57,9 +54,9 @@ def respond(
         top_p=top_p,
         pad_token_id=tokenizer.eos_token_id,
         do_sample=True,
-        attention_mask=attention_mask,  # Use the dynamically generated attention mask
+        attention_mask=attention_mask,
     )
-
+    
     # Decode the response
     response = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)